Update app.py (#30)

- Update app.py (f9cde6b6b17850035a6832df433db2ed755ae09c)

app.py

@@ -5,6 +5,10 @@ import inspect
 import pandas as pd
 from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, InferenceClientModel, WikipediaSearchTool, VisitWebpageTool, TransformersModel
 from huggingface_hub import login
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
 import os
 
 
@@ -18,35 +22,33 @@ token = os.getenv('HF_API_TOKEN')
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
     def __init__(self):
+        os.environ["HF_HOME"] = "/data/.huggingface"
         print("BasicAgent initialized.")
         # Ensure the token is not None
         if token is None:
             print("Hugging Face API token not found in environment variables.")
-        #login()
-        #model = HfApiModel(
-        #    max_tokens=2096,
-         #   temperature=0.5,
-          #  model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
-           # custom_role_conversions=None,
-        #)
-
-        #model = InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-1.5B-Instruct",
-         #                            max_tokens=2096,
-          #                           token=token,
-           #                          custom_role_conversions=None)
-
-        # Initialize the model with TransformersModel
-        model = TransformersModel(
-            model_id="gpt2",
-            max_new_tokens=4096,
+        
+        model_name = "gpt2"  # or "gpt2-medium", "gpt2-large"
+        
+        # Load tokenizer and inject a minimal chat_template
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if not getattr(tokenizer, "chat_template", None):
+            tokenizer.chat_template = (
+                "<|im_start|>system\n"
+                "{system}\n"
+                "<|im_end|>\n"
+                "<|im_start|>user\n"
+                "{user}\n"
+                "<|im_end|>\n"
+                "<|im_start|>assistant\n"
+            )
+        
+        # Load the GPT-2 model (FP16 if you like, but default is fine)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
             device_map="auto"
         )
-        self.agent = CodeAgent(
-            #model = HfApiModel(model_id="Qwen/Qwen2.5-Coder-7B-Instruct",token=token),
-            model = model,
-            tools = [DuckDuckGoSearchTool(),WikipediaSearchTool(),VisitWebpageTool()],
-            #additional_authorized_imports = ["pandas", "numpy"]
-        )
+
         SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question. 
         Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 
         YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
@@ -54,6 +56,13 @@ class BasicAgent:
         If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
         If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
         """
+        
+        self.agent = CodeAgent(
+            model = model,
+            tools = [DuckDuckGoSearchTool(),WikipediaSearchTool(),VisitWebpageTool()],
+            #additional_authorized_imports = ["pandas", "numpy"]
+        )
+
         self.agent.prompt_templates["system_prompt"] = self.agent.prompt_templates["system_prompt"] + SYSTEM_PROMPT
         
     def __call__(self, question: str) -> str: