Update app.py

app.py

@@ -24,43 +24,39 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 load_dotenv()
 
 import os
-import io
-import contextlib
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from smolagents import Tool
 
-class ZephyrToolCallingAgentTool(Tool):
-    name = "zephyr_tool_agent"
-    description = "Uses Zephyr-7B to answer questions using code or reasoning"
-
-    inputs = {
-        "question": {
-            "type": "string",
-            "description": "The user's question involving reasoning or code execution."
-        }
-    }
-    output_type = "string"
+class CodeLlamaToolCallingAgentTool:
+    name = "code_llama_tool_agent"
+    description = "Uses Code Llama to answer questions using code or reasoning"
 
     def __init__(self):
-        self.model_id = "HuggingFaceH4/zephyr-7b-beta"
-        token = os.getenv("HF_TOKEN")
+        self.model_id = "meta/code-llama-7b-instruct"
+        token = os.getenv("HF_TOKEN")  # Optional unless private
 
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, token=token)
         self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_id,
-            device_map="auto",
-            torch_dtype="auto",
-            token=token
+            self.model_id, device_map="auto", torch_dtype="auto", token=token
         )
+
         self.pipeline = pipeline(
             "text-generation",
             model=self.model,
             tokenizer=self.tokenizer,
-            max_new_tokens=128, # 128 or 512
+            max_new_tokens=512,
             temperature=0.2
         )
 
-    def forward(self, question: str) -> str:
+    def _run_code(self, code: str) -> str:
+        buffer = io.StringIO()
+        try:
+            with contextlib.redirect_stdout(buffer):
+                exec(code, {})
+            return buffer.getvalue().strip()
+        except Exception as e:
+            return f"Error during code execution: {e}"
+
+    def run(self, question: str) -> str:
         prompt = f"""You are a helpful assistant. Use code to solve questions that involve calculations.
 If code is needed, return a block like <tool>code</tool>. End your answer with <final>answer</final>.
 
@@ -69,9 +65,11 @@ Answer:"""
 
         result = self.pipeline(prompt)[0]["generated_text"]
 
+        # Process result
         if "<tool>" in result and "</tool>" in result:
             code = result.split("<tool>")[1].split("</tool>")[0].strip()
-            return self._run_code(code)
+            output = self._run_code(code)
+            return f"FINAL ANSWER (code output): {output}"
 
         elif "<final>" in result and "</final>" in result:
             final = result.split("<final>")[1].split("</final>")[0].strip()
@@ -79,15 +77,6 @@ Answer:"""
 
         return "Could not determine how to respond. No <tool> or <final> block detected."
 
-    def _run_code(self, code: str) -> str:
-        buffer = io.StringIO()
-        try:
-            with contextlib.redirect_stdout(buffer):
-                exec(code, {})
-            return f"FINAL ANSWER (code output): {buffer.getvalue().strip()}"
-        except Exception as e:
-            return f"Error during code execution: {e}"
-
 
 #from smolagents import Tool
 #from langchain_community.document_loaders import WikipediaLoader
@@ -262,7 +251,7 @@ class BasicAgent:
         video_transcription_tool = VideoTranscriptionTool()
 
         # ✅ New Mistral-based Tool
-        zephyr_tool = ZephyrToolCallingAgentTool()
+        my_tool = CodeLlamaToolCallingAgentTool()
 
         system_prompt = f"""
 You are my general AI assistant. Your task is to answer the question I asked.
@@ -281,7 +270,7 @@ If the answer is a comma-separated list, apply the above rules for each element
                 keywords_extract_tool, speech_to_text_tool,
                 visit_webpage_tool, final_answer_tool,
                 parse_excel_to_json, video_transcription_tool,
-                zephyr_tool  # 🔧 Add here
+                my_tool  # 🔧 Add here
             ],
             add_base_tools=True
         )