Update app.py

app.py

@@ -10,6 +10,7 @@ MAX_MODEL_TOKENS = 131072
 MAX_NEW_TOKENS = 4096
 MAX_CHUNK_TOKENS = 8192
 PROMPT_OVERHEAD = 300
+BATCH_SIZE = 3  # NEW: batching chunks to avoid overload
 
 # Paths
 persistent_dir = "/data/hf_cache"
@@ -83,14 +84,14 @@ def init_agent() -> TxAgent:
     agent.init_model()
     return agent
 
-# ⚡ Optimized parallel analysis
-def analyze_parallel(agent, chunks: List[str], max_workers: int = 6) -> List[str]:
-    results = [None] * len(chunks)
+# ⚡ Parallel analyzer with batching
+def analyze_parallel(agent, batch_chunks: List[List[str]], max_workers: int = 3) -> List[str]:
+    results = [None] * len(batch_chunks)
 
-    def worker(index, chunk):
-        prompt = build_prompt(chunk)
+    def worker(index, batch):
+        prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
         if estimate_tokens(prompt) > MAX_MODEL_TOKENS:
-            return f"❌ Chunk {index+1} too long. Skipped."
+            return f"❌ Batch {index+1} too long. Skipped."
         response = ""
         try:
             for r in agent.run_gradio_chat(
@@ -112,16 +113,16 @@ def analyze_parallel(agent, chunks: List[str], max_workers: int = 6) -> List[str
                     response += r.content
             return clean_response(response)
         except Exception as e:
-            return f"❌ Error in chunk {index+1}: {str(e)}"
+            return f"❌ Error in batch {index+1}: {str(e)}"
 
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        futures = {executor.submit(worker, idx, chunk): idx for idx, chunk in enumerate(chunks)}
+        futures = {executor.submit(worker, idx, batch): idx for idx, batch in enumerate(batch_chunks)}
         for future in futures:
             idx = futures[future]
             try:
                 results[idx] = future.result()
             except Exception as e:
-                results[idx] = f"❌ Error in chunk {idx+1}: {str(e)}"
+                results[idx] = f"❌ Error in batch {idx+1}: {str(e)}"
 
     gc.collect()
     return results
@@ -157,13 +158,14 @@ def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Di
     try:
         extracted = extract_text_from_excel(file.name)
         chunks = split_text(extracted)
-        messages.append({"role": "assistant", "content": f"🔍 Split into {len(chunks)} chunks. Analyzing..."})
+        batch_chunks = [chunks[i:i+BATCH_SIZE] for i in range(0, len(chunks), BATCH_SIZE)]
+        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batch_chunks)} batches. Analyzing..."})
 
-        chunk_results = analyze_parallel(agent, chunks, max_workers=6)
+        chunk_results = analyze_parallel(agent, batch_chunks, max_workers=3)
         valid = [res for res in chunk_results if not res.startswith("❌")]
 
         if not valid:
-            messages.append({"role": "assistant", "content": "❌ No valid chunk outputs."})
+            messages.append({"role": "assistant", "content": "❌ No valid batch outputs."})
             return messages, None
 
         summary = generate_final_summary(agent, "\n\n".join(valid))
@@ -179,21 +181,20 @@ def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Di
         messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
         return messages, None
 
+
 def create_ui(agent):
     with gr.Blocks(css="""
     html, body, .gradio-container {
         background-color: #0e1621;
         color: #e0e0e0;
         font-family: 'Inter', sans-serif;
-        padding: 0;
         margin: 0;
+        padding: 0;
     }
-
     h2, h3, h4 {
         color: #89b4fa;
         font-weight: 600;
     }
-
     button.gr-button-primary {
         background-color: #007bff !important;
         color: white !important;
@@ -203,26 +204,35 @@ def create_ui(agent):
         font-size: 16px !important;
         border: none;
     }
-
     button.gr-button-primary:hover {
         background-color: #0056b3 !important;
     }
-
     .gr-chatbot, .gr-markdown, .gr-file-upload {
         border-radius: 16px;
         background-color: #1b2533;
         border: 1px solid #2a2f45;
         padding: 10px;
     }
-
     .gr-chatbot .message {
         font-size: 16px;
         line-height: 1.6;
-        white-space: pre-wrap;
-        word-break: break-word;
+        padding: 12px 16px;
+        border-radius: 18px;
         margin: 8px 0;
+        max-width: 80%;
+        word-break: break-word;
+        white-space: pre-wrap;
+    }
+    .gr-chatbot .message.user {
+        background-color: #334155;
+        align-self: flex-end;
+        margin-left: auto;
+    }
+    .gr-chatbot .message.assistant {
+        background-color: #1e293b;
+        align-self: flex-start;
+        margin-right: auto;
     }
-
     .gr-file-upload .file-name {
         font-size: 14px;
         color: #89b4fa;
@@ -258,3 +268,4 @@ if __name__ == "__main__":
     except Exception as err:
         print(f"Startup failed: {err}")
         sys.exit(1)
+