Update app.py

app.py

@@ -6,33 +6,34 @@ import gradio as gr
 import pandas as pd
 from typing import Optional
 from dotenv import load_dotenv
-from transformers import pipeline
+from huggingface_hub import InferenceClient
 
-# ────────────────  Load Environment  ────────────────
+# ────────────────  Load Environment ────────────────
 load_dotenv()
 hf_token = os.getenv("HF_TOKEN")
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.1"
 
-# ────────────────  Load Falcon Model  ────────────────
+# ────────────────  Load Inference Client ────────────────
 try:
-    pipe = pipeline("text-generation", model="tiiuae/falcon-rw-1b", token=hf_token)
+    client = InferenceClient(model=MODEL_ID, token=hf_token)
 except Exception as e:
     raise RuntimeError(f"Model loading failed: {e}")
 
-# ────────────────  Agent Class  ────────────────
-class FalconAgent:
+# ────────────────  Agent Class ────────────────
+class MistralAgent:
     def __init__(self):
-        print("✅ FalconAgent initialized.")
+        print("✅ MistralAgent (Inference API) initialized.")
 
     def __call__(self, question: str) -> str:
         try:
-            prompt = f"{question}"
-            output = pipe(prompt, max_new_tokens=100, do_sample=False)
-            return output[0]["generated_text"].replace(prompt, "").strip()
+            prompt = f"[INST] {question.strip()} [/INST]"
+            response = client.text_generation(prompt, max_new_tokens=100, temperature=0.0, do_sample=False)
+            return response.strip()
         except Exception as e:
             return f"LLM Error: {e}"
 
-# ────────────────  API Helpers  ────────────────
+# ────────────────  API Helpers ────────────────
 def get_all_questions(api_url: str) -> list[dict]:
     resp = requests.get(f"{api_url}/questions", timeout=15)
     resp.raise_for_status()
@@ -48,8 +49,8 @@ def submit_answers(api_url: str, username: str, code_link: str, answers: list[di
     resp.raise_for_status()
     return resp.json()
 
-# ────────────────  Gradio Evaluation Logic  ────────────────
-def run_and_submit_all(profile: gr.OAuthProfile | None):
+# ────────────────  Gradio Evaluation Logic ────────────────
+def run_and_submit_all(profile: Optional[gr.OAuthProfile]):
     if not profile:
         return "❌ Please log in to Hugging Face using the button above.", None
     username = profile.username.strip()
@@ -58,7 +59,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     code_link = f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else ""
 
     try:
-        agent = FalconAgent()
+        agent = MistralAgent()
     except Exception as e:
         return f"❌ Error initializing agent: {e}", None
 
@@ -109,9 +110,9 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     )
     return final_status, pd.DataFrame(results_log)
 
-# ────────────────  Gradio UI  ────────────────
+# ────────────────  Gradio UI ────────────────
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Falcon-RW-1B Agent Evaluation")
+    gr.Markdown("# 🧠 Mistral-7B Agent Evaluation (via Inference API)")
     gr.Markdown(
         """
         **Instructions:**