Update app.py

app.py

@@ -15,31 +15,30 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
     def __init__(self):
-        print("🔄 Loading HF LLM pipeline...")
+        print("🔄 Loading Zephyr 7B...")
         self.llm = pipeline(
             "text-generation",
-            model="mistralai/Mistral-7B-Instruct-v0.2",
-            tokenizer="mistralai/Mistral-7B-Instruct-v0.2",
+            model="HuggingFaceH4/zephyr-7b-beta",
+            tokenizer="HuggingFaceH4/zephyr-7b-beta",
             max_new_tokens=200,
             temperature=0,
-            device=0 if torch.cuda.is_available() else -1  # Use GPU if available
+            device=0 if torch.cuda.is_available() else -1
         )
-        print("✅ LLM Loaded.")
+        print("✅ Zephyr 7B Loaded.")
 
     def __call__(self, question: str) -> str:
         print(f"🤖 Agent received question: {question}")
 
-        # Format as instruction
-        prompt = f"[INST] {question} [/INST]"
-        
+        prompt = f"<|system|>You are a helpful assistant.<|user|>{question}<|assistant|>"
+
         try:
             output = self.llm(prompt)[0]["generated_text"]
-            # Strip original prompt from output
-            answer = output.replace(prompt, "").strip()
+            # Post-process to remove prompt from output
+            answer = output.split("<|assistant|>")[-1].strip()
             print(f"✅ Answer: {answer}")
             return answer
         except Exception as e:
-            print(f"❌ Error: {e}")
+            print(f"❌ Error generating response: {e}")
             return "I don't know"
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):