Update app.py

app.py

@@ -3,6 +3,8 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -19,6 +21,30 @@ class BasicAgent:
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer
 
+class ZephyrAgent:
+    def __init__(self):
+        model_id = "HuggingFaceH4/zephyr-7b-beta"
+        print(f"Loading model: {model_id}")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.pipeline = pipeline(
+            "text-generation",
+            model=model_id,
+            tokenizer=self.tokenizer,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            device=0 if torch.cuda.is_available() else -1,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+        )
+        print("✅ ZephyrAgent initialized.")
+
+    def __call__(self, question: str) -> str:
+        prompt = f"<|system|>\nYou are a helpful AI assistant.\n<|user|>\n{question}\n<|assistant|>\n"
+        print(f"🧠 Prompting ZephyrAgent:\n{prompt[:100]}...")
+        response = self.pipeline(prompt, return_full_text=False)
+        generated_text = response[0]['generated_text'].strip()
+        return generated_text
+
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,
@@ -179,7 +205,8 @@ with gr.Blocks() as demo:
     test_button = gr.Button("Test Agent")
 
     def test_agent_response(question: str) -> str:
-        agent = BasicAgent()
+        # agent = BasicAgent()
+        agent = ZephyrAgent()
         return agent(question)
 
     test_button.click(fn=test_agent_response, inputs=question_input, outputs=answer_output)