Update app.py

app.py

@@ -9,7 +9,7 @@ import gc
 import re
 import torch
 from typing import List, Tuple, Dict
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # Directories
 persistent_dir = "/data/hf_cache"
@@ -100,18 +100,21 @@ def init_agent() -> TxAgent:
     agent.init_model()
     return agent
 
-def analyze_batches(agent, batches: List[List[str]]) -> List[str]:
+
+
+def analyze_batches(agent, batches: List[List[str]], max_workers: int = 3) -> List[str]:
     results = []
-    for batch in batches:
-        prompt = "\n\n".join(build_prompt(c) for c in batch)
+
+    def process_single_batch(batch):
+        prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
+        response = ""
         try:
-            response = ""
             for r in agent.run_gradio_chat(
                 message=prompt,
                 history=[],
                 temperature=0.0,
-                max_new_tokens=MAX_NEW_TOKENS,
-                max_token=MAX_MODEL_TOKENS,
+                max_new_tokens=4096,
+                max_token=131072,
                 call_agent=False,
                 conversation=[]
             ):
@@ -123,11 +126,15 @@ def analyze_batches(agent, batches: List[List[str]]) -> List[str]:
                             response += m.content
                 elif hasattr(r, "content"):
                     response += r.content
-            results.append(clean_response(response))
         except Exception as e:
-            results.append(f"❌ Error: {str(e)}")
-        torch.cuda.empty_cache()
-        gc.collect()
+            response = f"❌ Error: {str(e)}"
+        return clean_response(response)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [executor.submit(process_single_batch, batch) for batch in batches]
+        for future in as_completed(futures):
+            results.append(future.result())
+
     return results
 
 def generate_final_summary(agent, combined: str) -> str: