Update app.py

app.py

@@ -1,4 +1,5 @@
 import yaml
+import space
 import os
 from smolagents import GradioUI, CodeAgent, TransformersModel
 
@@ -10,12 +11,6 @@ from tools.get_current_weather import SimpleTool as GetCurrentWeather
 from tools.final_answer import FinalAnswerTool as FinalAnswer
 
 
-
-model = TransformersModel(
-max_new_tokens=1000,
-model_id='Qwen/Qwen2.5-Coder-3B-Instruct',
-)
-
 get_coordinates = GetCoordinates()
 get_current_weather = GetCurrentWeather()
 final_answer = FinalAnswer()
@@ -24,20 +19,30 @@ final_answer = FinalAnswer()
 with open(os.path.join(CURRENT_DIR, "prompts.yaml"), 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
-agent_get_weather_agent = CodeAgent(
-    model=model,
-    tools=[get_coordinates, get_current_weather],
-    managed_agents=[],
-    max_steps=20,
-    verbosity_level=2,
-    grammar=None,
-    planning_interval=None,
-    name='get_weather_agent',
-    description=None,
-    executor_type='local',
-    executor_kwargs={},
-    max_print_outputs_length=None,
-    prompt_templates=prompt_templates
-)
-if __name__ == "__main__":
+@spaces.GPU
+def call_agent():
+    model = TransformersModel(
+        max_new_tokens=1000,
+        model_id='Qwen/Qwen2.5-Coder-3B-Instruct',
+    )
+    
+    agent_get_weather_agent = CodeAgent(
+        model=model,
+        tools=[get_coordinates, get_current_weather],
+        managed_agents=[],
+        max_steps=20,
+        verbosity_level=2,
+        grammar=None,
+        planning_interval=None,
+        name='get_weather_agent',
+        description=None,
+        executor_type='local',
+        executor_kwargs={},
+        max_print_outputs_length=None,
+        prompt_templates=prompt_templates
+    )
     GradioUI(agent_get_weather_agent).launch()
+
+
+if __name__ == "__main__":
+    call_agent()