Update opensource_inference.py

opensource_inference.py

@@ -1,3 +1,5 @@
+%%writefile main.py
+import argparse
 import os
 import json
 import numpy as np
@@ -452,19 +454,57 @@ Conversation:
         }
 
 
-# Example usage
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Sales Conversion Predictor")
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="/content/sales-conversion-model-reinf-learning/sales_conversion_model",
+        help="Path to the trained PPO model zip file."
+    )
+    parser.add_argument(
+        "--embedding_model_name",
+        type=str,
+        default="BAAI/bge-m3", # Defaulting to bge-m3 as per example
+        help="Name of the Hugging Face embedding model (e.g., 'BAAI/bge-m3', 'BAAI/bge-large-en-v1.5')."
+    )
+    parser.add_argument(
+        "--llm_gguf_path",
+        type=str,
+        default="unsloth/gemma-3-4b-it-GGUF", # Defaulting to a repo ID as per example
+        help="Path to the GGUF LLM model file, a local directory containing GGUF files, or a HuggingFace repo_id."
+    )
+    parser.add_argument(
+        "--no_gpu",
+        action="store_true",
+        help="Disable GPU usage (use CPU only)."
+    )
+    parser.add_argument(
+        "--n_gpu_layers",
+        type=int,
+        default=-1, # Default to all layers on GPU for llama.cpp
+        help="Number of LLM layers to offload to GPU. -1 for all, 0 for none."
+    )
+    parser.add_argument(
+        "--n_ctx",
+        type=int,
+        default=2048,
+        help="Context window size for the LLM."
+    )
+
+    args = parser.parse_args()
+
     # Initialize predictor with GGUF model
     predictor = SalesConversionPredictor(
-        model_path="/content/sales-conversion-model-reinf-learning/sales_conversion_model", # path to the model
-        embedding_model_name="BAAI/bge-m3",
-        llm_gguf_path="unsloth/gemma-3-4b-it-GGUF",  # Update this path!
-        use_gpu=True,
-        n_gpu_layers=20,  # Use all layers on GPU
-        n_ctx=2048,      # Context window size
-        use_mini_embeddings=True  # Set to match how the model was trained
+        model_path=args.model_path,
+        embedding_model_name=args.embedding_model_name,
+        llm_gguf_path=args.llm_gguf_path,
+        use_gpu=not args.no_gpu,
+        n_gpu_layers=args.n_gpu_layers,
+        n_ctx=args.n_ctx,
+        use_mini_embeddings=True # Kept from original, PPO model should match this if it affects state vector.
+                                 # Currently, embedding dim is fixed at 1024 in code.
     )
-
     # Test with different conversation scenarios
     scenarios = [
         {