Update app.py

app.py

@@ -23,53 +23,67 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 #Load environment variables
 load_dotenv()
 
-import os
+import io
+import contextlib
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from smolagents import Tool, CodeAgent, DuckDuckGoSearchTool, HfApiModel, FinalAnswerTool
+
+class CodeLlamaTool(Tool):
+    name = "code_llama_tool"
+    description = "Uses Code Llama 7B Instruct to answer questions with code or reasoning"
 
-class CodeLlamaToolCallingAgentTool:
-    name = "code_llama_tool_agent"
-    description = "Uses Code Llama to answer questions using code or reasoning"
+    inputs = {
+        "question": {
+            "type": "string",
+            "description": "The user's question requiring reasoning or code execution."
+        }
+    }
+    output_type = "string"
 
     def __init__(self):
-        self.model_id = "meta/code-llama-7b-instruct"
-        token = os.getenv("HF_TOKEN")  # Optional unless private
+        self.model_id = "codellama/CodeLlama-7b-Instruct-hf"
+        token = os.getenv("HF_TOKEN")
 
         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.pipeline = pipeline(
             "text-generation",
             model=self.model,
             tokenizer=self.tokenizer,
             max_new_tokens=512,
-            temperature=0.1
+            temperature=0.2
         )
 
-    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>.
-
+    def forward(self, question: str) -> str:
+        prompt = f"""
+You are an advanced reasoning assistant. Use Python code if helpful.
+
+Instructions:
+- Solve step-by-step.
+- Wrap any code in <tool>...</tool>.
+- End your answer with <final>...</final> with the final result only.
+
+Example:
+Question: What is 3 times the square root of 49?
+Answer:
+<tool>
+import math
+print(3 * math.sqrt(49))
+</tool>
+<final>21.0</final>
+
+Now solve:
 Question: {question}
-Answer:"""
+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()
-            output = self._run_code(code)
-            return f"FINAL ANSWER (code output): {output}"
+            return self._run_code(code)
 
         elif "<final>" in result and "</final>" in result:
             final = result.split("<final>")[1].split("</final>")[0].strip()
@@ -77,6 +91,16 @@ 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
@@ -250,8 +274,8 @@ class BasicAgent:
         final_answer_tool = FinalAnswerTool()
         video_transcription_tool = VideoTranscriptionTool()
 
-        # ✅ New Mistral-based Tool
-        my_tool = CodeLlamaToolCallingAgentTool()
+        # ✅ New Llama Tool
+        my_tool = CodeLlamaTool()
 
         system_prompt = f"""
 You are my general AI assistant. Your task is to answer the question I asked.