Update app.py

app.py

@@ -10,45 +10,54 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-from transformers import pipeline
+# ---------- MODIFICATIONS BEGIN ----------
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import os
 import re
 
 class BasicAgent:
     def __init__(self):
-        print("Hybrid Agent Initialized: Regex + FLAN")
-        self.flan = pipeline("text2text-generation", model="google/flan-t5-base", device=-1)
+        print("Hybrid Agent with Mistral Model Initialized")
+
+        model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.getenv("HF_NEW_API_TOKEN"))
+        self.model = AutoModelForCausalLM.from_pretrained(model_id, token=os.getenv("HF_NEW_API_TOKEN"))
+        self.model.to("cpu")
+        self.model.eval()
 
     def classify(self, question: str) -> str:
         q = question.lower()
-        if any(k in q for k in ["youtube", ".mp3", ".wav", "image", ".png", ".jpg", "attached", "video"]):
+        if any(x in q for x in ["youtube", ".mp3", "image", "video", "attached", ".wav"]):
             return "media"
-        if any(k in q for k in ["|", "*", "subset", "commutative", "table", "="]):
+        if any(x in q for x in ["|", "*", "subset", "commutative", "table", "="]):
             return "logic"
-        return "nlp"
+        return "mistral"
 
-    def handle_media(self, question: str) -> str:
+    def handle_media(self, q: str) -> str:
         return "I'm unable to process audio, video, or file-based questions."
 
-    def handle_logic(self, question: str) -> str:
-        q = question.lower()
-
-        # Example: subset + commutative detection
+    def handle_logic(self, q: str) -> str:
+        q = q.lower()
         if "not commutative" in q and "subset" in q:
             return "a,b,c"
-
-        # More logic can be added here using regex
-        match = re.search(r"what country had the least number of athletes.*1928", q)
-        if match:
-            return "AFG"  # For example, return Afghanistan if it fits
-
-        return "I couldn't process this structured logic. Please verify manually."
-
-    def handle_nlp(self, question: str) -> str:
-        try:
-            result = self.flan(f"Answer this clearly and briefly:\n{question.strip()}", max_new_tokens=256)
-            return result[0]["generated_text"].strip()
-        except Exception as e:
-            return f"❌ FLAN Error: {e}"
+        return "I couldn't solve this logic-based question."
+
+    def handle_mistral(self, question: str) -> str:
+        prompt = f"<s>[INST] {question.strip()} [/INST]"
+        inputs = self.tokenizer(prompt, return_tensors="pt").to("cpu")
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=256,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95
+            )
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response.split("[/INST]")[-1].strip()
 
     def __call__(self, question: str) -> str:
         qtype = self.classify(question)
@@ -59,7 +68,8 @@ class BasicAgent:
         elif qtype == "logic":
             return self.handle_logic(question)
         else:
-            return self.handle_nlp(question)
+            return self.handle_mistral(question)
+# ---------- MODIFICATIONS END ----------
             
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """