Update app.py

app.py

@@ -33,26 +33,30 @@ MAX_CHUNK_TOKENS = 8192
 MAX_NEW_TOKENS = 2048
 PROMPT_OVERHEAD = 500
 
+
 def clean_response(text: str) -> str:
     text = re.sub(r"\[.*?\]|\bNone\b", "", text, flags=re.DOTALL)
     text = re.sub(r"\n{3,}", "\n\n", text)
     text = re.sub(r"[^\n#\-\*\w\s\.,:\(\)]+", "", text)
     return text.strip()
 
+
 def estimate_tokens(text: str) -> int:
     return len(text) // 3.5 + 1
 
+
 def extract_text_from_excel(file_path: str) -> str:
     all_text = []
     xls = pd.ExcelFile(file_path)
     for sheet_name in xls.sheet_names:
         df = xls.parse(sheet_name).astype(str).fillna("")
-        rows = df.apply(lambda row: " | ".join(row), axis=1)
-        sheet_text = [f"[{sheet_name}] {line}" for line in rows]
+        rows = df.apply(lambda row: " | ".join([cell for cell in row if cell.strip()]), axis=1)
+        sheet_text = [f"[{sheet_name}] {line}" for line in rows if line.strip()]
         all_text.extend(sheet_text)
     return "\n".join(all_text)
 
-def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
+
+def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS, max_chunks: int = 30) -> List[str]:
     effective_max = max_tokens - PROMPT_OVERHEAD
     lines, chunks, curr_chunk, curr_tokens = text.split("\n"), [], [], 0
     for line in lines:
@@ -60,14 +64,17 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> Lis
         if curr_tokens + t > effective_max:
             if curr_chunk:
                 chunks.append("\n".join(curr_chunk))
+            if len(chunks) >= max_chunks:
+                break
             curr_chunk, curr_tokens = [line], t
         else:
             curr_chunk.append(line)
             curr_tokens += t
-    if curr_chunk:
+    if curr_chunk and len(chunks) < max_chunks:
         chunks.append("\n".join(curr_chunk))
     return chunks
 
+
 def build_prompt_from_text(chunk: str) -> str:
     return f"""
 ### Unstructured Clinical Records
@@ -87,6 +94,7 @@ Analyze the following clinical notes and provide a detailed, concise summary foc
 Respond in well-structured bullet points with medical reasoning.
 """
 
+
 def init_agent():
     tool_path = os.path.join(tool_cache_dir, "new_tool.json")
     if not os.path.exists(tool_path):
@@ -103,6 +111,7 @@ def init_agent():
     agent.init_model()
     return agent
 
+
 def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
     messages = chatbot_state if chatbot_state else []
     if file is None or not hasattr(file, "name"):
@@ -116,7 +125,11 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
     def analyze_chunk(i, chunk):
         prompt = build_prompt_from_text(chunk)
         response = ""
-        for res in agent.run_gradio_chat(message=prompt, history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[]):
+        for res in agent.run_gradio_chat(
+            message=prompt, history=[], temperature=0.2,
+            max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS,
+            call_agent=False, conversation=[]
+        ):
             if isinstance(res, str):
                 response += res
             elif hasattr(res, "content"):
@@ -127,7 +140,7 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
                         response += r.content
         return i, clean_response(response)
 
-    with ThreadPoolExecutor(max_workers=1) as executor:
+    with ThreadPoolExecutor(max_workers=4) as executor:
         futures = [executor.submit(analyze_chunk, i, c) for i, c in enumerate(chunks)]
         for f in as_completed(futures):
             i, result = f.result()
@@ -141,7 +154,11 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
     messages.append({"role": "assistant", "content": "📊 Generating final report..."})
 
     final_report = ""
-    for res in agent.run_gradio_chat(message=summary_prompt, history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[]):
+    for res in agent.run_gradio_chat(
+        message=summary_prompt, history=[], temperature=0.2,
+        max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS,
+        call_agent=False, conversation=[]
+    ):
         if isinstance(res, str):
             final_report += res
         elif hasattr(res, "content"):
@@ -156,20 +173,22 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
     messages.append({"role": "assistant", "content": f"✅ Report generated and saved: {os.path.basename(report_path)}"})
     return messages, report_path
 
+
 def create_ui(agent):
     with gr.Blocks(css="""
         body {
             background: #10141f;
             color: #ffffff;
             font-family: 'Inter', sans-serif;
+            margin: 0;
+            padding: 0;
         }
         .gradio-container {
             padding: 30px;
-            max-width: 900px;
-            margin: auto;
-            border-radius: 16px;
+            width: 100vw;
+            max-width: 100%;
+            border-radius: 0;
             background-color: #1a1f2e;
-            box-shadow: 0 0 20px rgba(0, 0, 0, 0.5);
         }
         .chatbot {
             background-color: #131720;
@@ -203,7 +222,8 @@ Upload clinical Excel records below and click **Analyze** to generate a medical
 
         def update_ui(file, current_state):
             messages, report_path = process_final_report(agent, file, current_state)
-            return messages, gr.update(visible=report_path is not None, value=report_path), messages
+            chat_format = [(msg["role"], msg["content"]) for msg in messages if isinstance(msg, dict)]
+            return chat_format, gr.update(visible=report_path is not None, value=report_path), messages
 
         analyze_btn.click(
             fn=update_ui,
@@ -213,6 +233,7 @@ Upload clinical Excel records below and click **Analyze** to generate a medical
 
     return demo
 
+
 if __name__ == "__main__":
     try:
         agent = init_agent()