Update app.py

app.py

@@ -10,43 +10,49 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-import requests
-import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 class BasicAgent:
     def __init__(self):
-        print("Mistral Agent using Inference API initialized.")
-        self.token = os.getenv("HF_NEW_API_TOKEN")
-        self.api_url = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
-        self.headers = {
-            "Authorization": f"Bearer {self.token}",
-            "Content-Type": "application/json"
-        }
+        print("Loading Mistral with manual generate()...")
+
+        model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+
+        # Load tokenizer and model (gated model → needs HF token access if private)
+        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"))
+
+        # CPU-only
+        self.model.to("cpu")
+        self.model.eval()
 
     def __call__(self, question: str) -> str:
-        print(f"Sending question to API: {question[:50]}...")
         prompt = f"<s>[INST] {question.strip()} [/INST]"
 
         try:
-            response = requests.post(
-                self.api_url,
-                headers=self.headers,
-                json={"inputs": prompt, "parameters": {"max_new_tokens": 256, "temperature": 0.7}},
-                timeout=60
-            )
-            response.raise_for_status()
-            output = response.json()
-
-            # Handle potential format differences
-            if isinstance(output, list) and "generated_text" in output[0]:
-                return output[0]["generated_text"].split("[/INST]")[-1].strip()
-            else:
-                print(f"⚠️ Unexpected response: {output}")
-                return "⚠️ Mistral returned an unexpected format."
+            # Tokenize the prompt
+            inputs = self.tokenizer(prompt, return_tensors="pt")
+            input_ids = inputs["input_ids"].to("cpu")
+
+            # Generate text
+            with torch.no_grad():
+                generated_ids = self.model.generate(
+                    input_ids,
+                    max_new_tokens=256,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.95
+                )
+
+            # Decode output
+            output = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            answer = output.split("[/INST]")[-1].strip()
+            return answer
 
         except Exception as e:
-            print(f"❌ Error during API call: {e}")
-            return f"❌ API Error: {str(e)}"
+            print(f"❌ Error during generation: {e}")
+            return f"❌ Model Error: {str(e)}"
             
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """