Update app.py

app.py

@@ -5,6 +5,7 @@ import torch
 import pandas as pd
 import inspect
 import os
+import gc
 
 import torch, transformers
 print(f"PyTorch: {torch.__version__}")
@@ -84,6 +85,7 @@ class BasicAgent:
                 self.agent = CodeAgent(
                     tools=[self.search_tool],
                     model=self.model,
+                    max_steps=30,
                     additional_authorized_imports=[
                         'math', 'statistics', 're',  # Basic computation
                         'requests', 'json',           # Web requests and JSON
@@ -117,7 +119,8 @@ class BasicAgent:
             print(f"🔧 Available tools: {[tool.__class__.__name__ for tool in self.agent.tools]}")
             
             # Run the agent
-            result = self.agent.run(formatted_question)
+            with torch.no_grad():
+                result = self.agent.run(formatted_question)
             
             print(f"Raw result: {result}")
 
@@ -148,6 +151,14 @@ class BasicAgent:
             return self._fallback_implementation(question)
 
 
+def cleanup_memory():
+    """Centralized memory cleanup function"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+    gc.collect()
+
+
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,
@@ -203,6 +214,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     answers_payload = []
     print(f"Running agent on {len(questions_data)} questions...")
     for item in questions_data:
+        cleanup_memory()
         task_id = item.get("task_id")
         question_text = item.get("question")
         if not task_id or question_text is None: