Update app.py

app.py

@@ -14,32 +14,40 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 class BasicAgent:
     def __init__(self):
-        print("Mistral Local Agent initialized.")
+        print("Zephyr Local Agent initialized.")
 
-        model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+        model_id = "HuggingFaceH4/zephyr-7b-beta"
 
-        # Load model and tokenizer directly
+        # Load model and tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
         self.model = AutoModelForCausalLM.from_pretrained(model_id)
 
-        # Create inference pipeline
+        # Create generation pipeline
         self.pipeline = pipeline(
             "text-generation",
             model=self.model,
             tokenizer=self.tokenizer,
-            device=-1
+            device=-1  # CPU
         )
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
 
         try:
-            prompt = f"<s>[INST] {question.strip()} [/INST]"
-            output = self.pipeline(prompt, max_new_tokens=256, temperature=0.7)
-
-            # Extract and clean the response
+            # Format for chat-style prompt
+            prompt = f"<|system|>You are a helpful assistant.<|user|>{question.strip()}<|assistant|>"
+
+            output = self.pipeline(
+                prompt,
+                max_new_tokens=256,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95
+            )
+
+            # Return model's response only (strip off prompt)
             generated_text = output[0]["generated_text"]
-            response = generated_text.split("[/INST]")[-1].strip()
+            response = generated_text.split("<|assistant|>")[-1].strip()
             return response
 
         except Exception as e: