Update app.py

app.py

@@ -6,7 +6,9 @@ import subprocess
 import pandas as pd
 import torch, spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from smolagents import CodeAgent, HfApiModel
+# from smolagents import CodeAgent, HfApiModel
+from smolagents.agent import LocalLLM, CodeAgent
+from smolagents.message import MessageRole
 from huggingface_hub import InferenceClient, hf_hub_download
 
 subprocess.run(["playwright", "install"], check=True)
@@ -15,40 +17,28 @@ subprocess.run(["playwright", "install"], check=True)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-class LocalLLM:
-    def __init__(self, pipe):
-        self.pipe = pipe
+# class LocalLLM:
+#     def __init__(self, pipe):
+#         self.pipe = pipe
 
-    def generate(self, prompt, **kwargs):
-        unsupported_keys = ["stop_sequences"] # Remove keys not accepted by HF pipelines
-        cleaned_kwargs = {k: v for k, v in kwargs.items() if k not in unsupported_keys}
-        # print(f"🧪 kwargs cleaned: {cleaned_kwargs.keys()}")
-        try:
-            outputs = self.pipe(prompt, **cleaned_kwargs)
-            # print(f"🧪 Raw output from pipe: {outputs}")
-            if isinstance(outputs, list) and isinstance(outputs[0], dict):
-                out = outputs[0]["generated_text"]
-            elif isinstance(outputs, list):
-                out = outputs[0]  # fallback if it's just a list of strings
-            else:
-                out = str(outputs)
-            print("🧪 Final object to return:", type(out), out[:100])
-            return {'role': 'assistant', 'content': [{'type':'text', 'text': out}]}
-        except Exception as e:
-            print(f"❌ Error in LocalLLM.generate(): {e}")
-            raise
-
-@spaces.GPU
-def load_llm(hf_token):
-    model_id = "meta-llama/Llama-3.1-8B-Instruct"
-    tok = AutoTokenizer.from_pretrained(model_id, token=hf_token)
-    mod = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        device_map="auto",  # auto-distributes to GPU
-        token=hf_token
-    )
-    return pipeline("text-generation", model=mod, tokenizer=tok, max_new_tokens=512)
+#     def generate(self, prompt, **kwargs):
+#         unsupported_keys = ["stop_sequences"] # Remove keys not accepted by HF pipelines
+#         cleaned_kwargs = {k: v for k, v in kwargs.items() if k not in unsupported_keys}
+#         # print(f"🧪 kwargs cleaned: {cleaned_kwargs.keys()}")
+#         try:
+#             outputs = self.pipe(prompt, **cleaned_kwargs)
+#             # print(f"🧪 Raw output from pipe: {outputs}")
+#             if isinstance(outputs, list) and isinstance(outputs[0], dict):
+#                 out = outputs[0]["generated_text"]
+#             elif isinstance(outputs, list):
+#                 out = outputs[0]  # fallback if it's just a list of strings
+#             else:
+#                 out = str(outputs)
+#             print("🧪 Final object to return:", type(out), out[:100])
+#             return {'role': 'assistant', 'content': [{'type':'text', 'text': out}]}
+#         except Exception as e:
+#             print(f"❌ Error in LocalLLM.generate(): {e}")
+#             raise
 
 def check_token_access():
     token = os.environ.get("HF_TOKEN", "")
@@ -71,10 +61,10 @@ def check_token_access():
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
-class BasicAgent:
-    def __init__(self):
+class BasicAgent(LocalLLM):
+    def __init__(self, model_id="meta-llama/Llama-3.1-8B-Instruct", hf_token=None):
         print("BasicAgent initialized.")
-        print("ENV-HF_TOKEN-LEN", len(os.environ["HF_TOKEN"]), file=sys.stderr)
+        print("ENV-HF_TOKEN-LEN", len(hf_token), file=sys.stderr)
         check_token_access()
 
         # Local test
@@ -96,9 +86,22 @@ class BasicAgent:
         # See list of base tools in https://github.com/huggingface/smolagents/blob/main/src/smolagents/default_tools.py
         
         # Initialize the agent
-        pipe = load_llm(hf_token = os.environ["HF_TOKEN"])
-        self.model = LocalLLM(pipe)
-        self.agent = CodeAgent(model=self.model, tools=[], add_base_tools=True)
+        tok = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+        mod = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16,
+            device_map="auto",  # auto-distributes to GPU
+            token=hf_token
+        )
+        self.pipe = pipeline(
+            "text-generation",
+            model=mod,
+            tokenizer=tok,
+            max_new_tokens=512,
+            temperature=0.0,
+        )
+        
+        self.agent = CodeAgent(model=self, tools=[], add_base_tools=True)
         
         # self.agent = CodeAgent(
         #     model=model, 
@@ -106,11 +109,34 @@ class BasicAgent:
         #     add_base_tools=True
         # )
         
+    def _serialize_messages(self, messages):
+    prompt = []
+    for m in messages:
+        role = m['role'].value  # "system" / "user" / "assistant"
+        text = "".join([c['text'] for c in m['content']])
+        prompt.append(f"{role}: {text}")
+    return "\n".join(prompt)
+
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        response = self.agent.run(question)
+
+        allowed = {"max_new_tokens", "temperature", "top_k", "top_p"}
+        gen_kwargs = {k: v for k, v in kwargs.items() if k in allowed}
+        prompt_str = (
+            self._serialize_messages(prompt)
+            if isinstance(prompt, list)
+            else prompt
+        )
+        outputs = self.pipe(prompt_str, **gen_kwargs)
+        response = outputs[0]["generated_text"]
+        # response = self.agent.run(question)
         print(f"Agent returning its generated answer: {response}")
-        return response
+        
+        # wrap back into a chat message dict
+        return {
+            "role": MessageRole.ASSISTANT,
+            "content": [{"type": "text", "text": response}],
+        }
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
@@ -119,6 +145,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     # --- Determine HF Space Runtime URL and Repo URL ---
     space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code
+    hf_token = os.getenv("HF_TOKEN")
 
     if profile:
         username= f"{profile.username}"
@@ -133,7 +160,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = BasicAgent()
+        agent = BasicAgent(hf_token=hf_token).agent
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -173,7 +200,8 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
-            submitted_answer = agent(question_text)
+            msg = agent.run(question_text)
+            submitted_answer = msg["content"][0]["text"]
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e:
@@ -269,6 +297,7 @@ if __name__ == "__main__":
     # Check for SPACE_HOST and SPACE_ID at startup for information
     space_host_startup = os.getenv("SPACE_HOST")
     space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup
+    hf_token = os.getenv("HF_TOKEN")
 
     if space_host_startup:
         print(f"✅ SPACE_HOST found: {space_host_startup}")
@@ -286,8 +315,8 @@ if __name__ == "__main__":
     print("-"*(60 + len(" App Starting ")) + "\n")
 
     # Test the agent
-    agent = BasicAgent()
-    agent.agent.run("What is 2+2?")
+    agent = BasicAgent(hf_token=hf_token).agent
+    agent.run("What is 2+2?")
 
     print("Launching Gradio Interface for Basic Agent Evaluation...")
     demo.launch(debug=True, share=False)