Update myagent.py

myagent.py

@@ -1,59 +1,75 @@
-import os
-from smolagents import CodeAgent, ToolCallingAgent
-from smolagents import OpenAIServerModel
-from tools.fetch import fetch_webpage
-from tools.yttranscript import get_youtube_transcript, get_youtube_title_description
-import myprompts
-
-# --- Basic Agent Definition ---
-class BasicAgent:
-    def __init__(self):
-        print("BasicAgent initialized.")
-    def __call__(self, question: str) -> str:
-
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-
-        try:
-            # Use the reviewer agent to determine if the question can be answered by a model or requires code
-            print("Calling reviewer agent...")
-            reviewer_answer = reviewer_agent.run(myprompts.review_prompt + "\nThe question is:\n" + question)
-            print(f"Reviewer agent answer: {reviewer_answer}")
-
-            question = question + '\n' + myprompts.output_format
-            fixed_answer = ""
-
-            if reviewer_answer == "code":
-                fixed_answer = gaia_agent.run(question)
-                print(f"Code agent answer: {fixed_answer}")
-                
-            elif reviewer_answer == "model":    
-                # If the reviewer agent suggests using the model, we can proceed with the model agent
-                print("Using model agent to answer the question.")
-                fixed_answer = model_agent.run(myprompts.model_prompt + "\nThe question is:\n" + question)
-                print(f"Model agent answer: {fixed_answer}")
-
-            return fixed_answer
-        except Exception as e:
-            error = f"An error occurred while processing the question: {e}"
-            print(error)
-            return error
-
-        
-      
-    
-model = OpenAIServerModel(
-    model_id="gpt-4.1-nano",
-    api_base="https://api.openai.com/v1",
-    api_key=os.environ["OPENAI_API_KEY"],
-)
-
-reviewer_agent= ToolCallingAgent(model=model, tools=[])
-model_agent = ToolCallingAgent(model=model,tools=[fetch_webpage])
-gaia_agent = CodeAgent(tools=[fetch_webpage,get_youtube_title_description,get_youtube_transcript ], model=model)
-
-if __name__ == "__main__":
-    # Example usage
-    question = "What was the actual enrollment of the Malko competition in 2023?"
-    agent = BasicAgent()
-    answer = agent(question)
+import os
+from smolagents import CodeAgent, ToolCallingAgent
+from smolagents import OpenAIServerModel
+from tools.fetch import fetch_webpage
+from tools.yttranscript import get_youtube_transcript, get_youtube_title_description
+import myprompts
+
+# --- Basic Agent Definition ---
+class BasicAgent:
+    def __init__(self):
+        print("BasicAgent initialized.")
+    def __call__(self, question: str) -> str:
+
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+
+        try:
+            # Use the reviewer agent to determine if the question can be answered by a model or requires code
+            print("Calling reviewer agent...")
+            reviewer_answer = reviewer_agent.run(myprompts.review_prompt + "\nThe question is:\n" + question)
+            print(f"Reviewer agent answer: {reviewer_answer}")
+
+            question = question + '\n' + myprompts.output_format
+            fixed_answer = ""
+
+            if reviewer_answer == "code":
+                fixed_answer = gaia_agent.run(question)
+                print(f"Code agent answer: {fixed_answer}")
+                
+            elif reviewer_answer == "model":    
+                # If the reviewer agent suggests using the model, we can proceed with the model agent
+                print("Using model agent to answer the question.")
+                fixed_answer = model_agent.run(myprompts.model_prompt + "\nThe question is:\n" + question)
+                print(f"Model agent answer: {fixed_answer}")
+
+            return fixed_answer
+        except Exception as e:
+            error = f"An error occurred while processing the question: {e}"
+            print(error)
+            return error
+
+        
+      
+    
+# model = OpenAIServerModel(
+#     model_id="gpt-4.1-nano",
+#     api_base="https://api.openai.com/v1",
+#     api_key=os.environ["OPENAI_API_KEY"],
+# )
+
+
+MODEL_NAME = "meta-llama/Llama-3.2-3B"  # 3B isn't released by Meta officially, but use 8B or a 3B variant like TinyLlama if needed
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model_init = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",
+    torch_dtype=torch.float16  # or bfloat16
+)
+
+def model(prompt: str, max_new_tokens=512):
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+    output_ids = model.generate(input_ids, max_new_tokens=max_new_tokens)
+    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return output
+
+reviewer_agent= ToolCallingAgent(model=model, tools=[])
+model_agent = ToolCallingAgent(model=model,tools=[fetch_webpage])
+gaia_agent = CodeAgent(tools=[fetch_webpage,get_youtube_title_description,get_youtube_transcript ], model=model)
+
+if __name__ == "__main__":
+    # Example usage
+    question = "What was the actual enrollment of the Malko competition in 2023?"
+    agent = BasicAgent()
+    answer = agent(question)
     print(f"Answer: {answer}")