Update app.py

app.py

@@ -4,7 +4,7 @@ import pandas as pd
 import json
 import gradio as gr
 from typing import List
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import ThreadPoolExecutor
 import hashlib
 import shutil
 import re
@@ -27,7 +27,7 @@ except ImportError:
     HAS_PYPDFIUM2 = False
     import pdfplumber
 
-# Configure logging and suppress warnings
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logging.getLogger("pdfminer").setLevel(logging.ERROR)
 logger = logging.getLogger(__name__)
@@ -67,76 +67,58 @@ def file_hash(path: str) -> str:
     with open(path, "rb") as f:
         return hashlib.md5(f.read()).hexdigest()
 
-async def extract_all_pages_async(file_path: str, progress_callback=None, use_ocr=False) -> str:
+async def extract_all_pages_async(file_path: str, progress_callback=None, force_ocr=False) -> str:
     try:
+        extracted_text = ""
+        total_pages = 0
+        text_chunks = []
+
         if HAS_PYPDFIUM2:
             pdf = pdfium.PdfDocument(file_path)
             total_pages = len(pdf)
             if total_pages == 0:
                 return ""
 
-            batch_size = 5
-            batches = [(i, min(i + batch_size, total_pages)) for i in range(0, total_pages, batch_size)]
-            text_chunks = [""] * total_pages
-            processed_pages = 0
-
-            def extract_batch(start: int, end: int) -> List[tuple]:
-                results = []
-                for i in range(start, end):
-                    page = pdf[i]
-                    text = page.get_textpage().get_text_range() or ""
-                    if not text.strip() and use_ocr and 'pytesseract' in sys.modules:
-                        bitmap = page.render(scale=2).to_pil()
-                        text = pytesseract.image_to_string(bitmap, lang="eng")
-                    results.append((i, f"=== Page {i + 1} ===\n{text.strip()}"))
-                return results
-
-            loop = asyncio.get_event_loop()
+            def extract_page(i):
+                page = pdf[i]
+                text = page.get_textpage().get_text_range() or ""
+                if (not text.strip() or len(text) < 100) and force_ocr and 'pytesseract' in sys.modules:
+                    logger.info("Falling back to OCR for page %d", i + 1)
+                    bitmap = page.render(scale=2).to_pil()
+                    text = pytesseract.image_to_string(bitmap, lang="eng")
+                return (i, f"=== Page {i + 1} ===\n{text.strip()}")
+
             with ThreadPoolExecutor(max_workers=4) as executor:
-                futures = [loop.run_in_executor(executor, extract_batch, start, end) for start, end in batches]
-                for future in await asyncio.gather(*futures):
-                    for page_num, text in future:
-                        text_chunks[page_num] = text
-                        logger.debug("Page %d extracted: %s...", page_num + 1, text[:50])
-                    processed_pages += batch_size
+                futures = [executor.submit(extract_page, i) for i in range(total_pages)]
+                for future in as_completed(futures):
+                    page_num, text = future.result()
+                    text_chunks.append((page_num, text))
+                    logger.debug("Page %d extracted: %s...", page_num + 1, text[:50])
                     if progress_callback:
-                        progress_callback(min(processed_pages, total_pages), total_pages)
+                        progress_callback(page_num + 1, total_pages)
 
+            text_chunks.sort(key=lambda x: x[0])
+            extracted_text = "\n\n".join(chunk[1] for chunk in text_chunks if chunk[1].strip())
             pdf.close()
         else:
-            # Fallback to pdfplumber
             with pdfplumber.open(file_path) as pdf:
                 total_pages = len(pdf.pages)
                 if total_pages == 0:
                     return ""
 
-                batch_size = 5
-                batches = [(i, min(i + batch_size, total_pages)) for i in range(0, total_pages, batch_size)]
-                text_chunks = [""] * total_pages
-                processed_pages = 0
-
-                def extract_batch(start: int, end: int) -> List[tuple]:
-                    results = []
-                    with pdfplumber.open(file_path) as pdf:
-                        for i in range(start, end):
-                            page = pdf.pages[i]
-                            text = page.extract_text() or ""
-                            results.append((i, f"=== Page {i + 1} ===\n{text.strip()}"))
-                    return results
-
-                loop = asyncio.get_event_loop()
-                with ThreadPoolExecutor(max_workers=4) as executor:
-                    futures = [loop.run_in_executor(executor, extract_batch, start, end) for start, end in batches]
-                    for future in await asyncio.gather(*futures):
-                        for page_num, text in future:
-                            text_chunks[page_num] = text
-                            logger.debug("Page %d extracted: %s...", page_num + 1, text[:50])
-                        processed_pages += batch_size
-                        if progress_callback:
-                            progress_callback(min(processed_pages, total_pages), total_pages)
-
-        extracted_text = "\n\n".join(filter(None, text_chunks))
+                for i, page in enumerate(pdf.pages):
+                    text = page.extract_text() or ""
+                    text_chunks.append((i, f"=== Page {i + 1} ===\n{text.strip()}"))
+                    logger.debug("Page %d extracted: %s...", i + 1, text[:50])
+                    if progress_callback:
+                        progress_callback(i + 1, total_pages)
+
+            extracted_text = "\n\n".join(chunk[1] for chunk in text_chunks if chunk[1].strip())
+
         logger.info("Extracted %d pages, total length: %d chars", total_pages, len(extracted_text))
+        if len(extracted_text) < 1000 and not force_ocr and HAS_PYPDFIUM2 and 'pytesseract' in sys.modules:
+            logger.info("Text too short, retrying with OCR")
+            return await extract_all_pages_async(file_path, progress_callback, force_ocr=True)
         return extracted_text
     except Exception as e:
         logger.error("PDF processing error: %s", e)
@@ -151,10 +133,7 @@ def convert_file_to_json(file_path: str, file_type: str, progress_callback=None)
             return cache[cache_key]
 
         if file_type == "pdf":
-            text = asyncio.run(extract_all_pages_async(file_path, progress_callback, use_ocr=False))
-            if not text.strip() or "PDF processing error" in text:
-                logger.info("Retrying extraction with OCR for %s", file_path)
-                text = asyncio.run(extract_all_pages_async(file_path, progress_callback, use_ocr=True))
+            text = asyncio.run(extract_all_pages_async(file_path, progress_callback, force_ocr=False))
             result = json.dumps({"filename": os.path.basename(file_path), "content": text, "status": "initial"})
         elif file_type == "csv":
             df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str,
@@ -199,10 +178,12 @@ def clean_response(text: str) -> str:
     text = text.replace("\n\n\n", "\n\n")
     sections = {}
     current_section = None
+    seen_lines = set()
     for line in text.splitlines():
         line = line.strip()
-        if not line:
+        if not line or line in seen_lines:
             continue
+        seen_lines.add(line)
         section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line)
         if section_match:
             current_section = section_match.group(1)
@@ -213,13 +194,21 @@ def clean_response(text: str) -> str:
     cleaned = [f"### {heading}\n" + "\n".join(findings) for heading, findings in sections.items() if findings]
     result = "\n\n".join(cleaned).strip()
     logger.debug("Cleaned response length: %d chars", len(result))
-    return result or "No issues identified"
-
-def summarize_findings(combined_response: str) -> str:
-    if not combined_response or all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")):
-        return "### Summary of Clinical Oversights\nNo critical oversights identified in the provided records."
-    sections = {}
+    return result or "No oversights identified"
+
+def summarize_findings(all_responses: List[str]) -> str:
+    combined_response = "\n\n".join(all_responses)
+    if not combined_response or all("No oversights identified" in resp.lower() for resp in all_responses):
+        return "### Comprehensive Clinical Oversight Summary\nNo critical oversights were identified across the provided patient records after thorough analysis."
+
+    sections = {
+        "Missed Diagnoses": [],
+        "Medication Conflicts": [],
+        "Incomplete Assessments": [],
+        "Urgent Follow-up": []
+    }
     current_section = None
+    seen_findings = set()
     for line in combined_response.splitlines():
         line = line.strip()
         if not line:
@@ -227,15 +216,19 @@ def summarize_findings(combined_response: str) -> str:
         section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line)
         if section_match:
             current_section = section_match.group(1)
-            sections.setdefault(current_section, [])
             continue
-        if current_section and line.startswith("- "):
-            sections[current_section].append(line[2:])
-    summary_lines = [
-        f"- **{heading}**: {'; '.join(findings[:1])}. Risks: potential adverse outcomes. Recommend: urgent review."
-        for heading, findings in sections.items() if findings
-    ]
-    result = "### Summary of Clinical Oversights\n" + "\n".join(summary_lines) if summary_lines else "### Summary of Clinical Oversights\nNo critical oversights identified."
+        if current_section and line.startswith("- ") and line not in seen_findings:
+            sections[current_section].append(line)
+            seen_findings.add(line)
+
+    summary_lines = []
+    for heading, findings in sections.items():
+        if findings:
+            summary_lines.append(f"### {heading}")
+            for finding in findings:
+                summary_lines.append(f"{finding}\n  - **Risks**: Potential adverse outcomes if not addressed.\n  - **Recommendation**: Immediate clinical review and follow-up.")
+
+    result = "### Comprehensive Clinical Oversight Summary\n" + "\n".join(summary_lines) if summary_lines else "### Comprehensive Clinical Oversight Summary\nNo critical oversights identified."
     logger.debug("Summary length: %d chars", len(result))
     return result
 
@@ -267,8 +260,8 @@ def init_agent():
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
-        chatbot = gr.Chatbot(label="Detailed Analysis", height=600, type="messages")
-        final_summary = gr.Markdown(label="Summary of Clinical Oversights")
+        chatbot = gr.Chatbot(label="Detailed Analysis", height=600, type="messages", visible=False)
+        final_summary = gr.Markdown(label="Comprehensive Clinical Oversight Summary")
         file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
         msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False)
         send_btn = gr.Button("Analyze", variant="primary")
@@ -276,9 +269,9 @@ def create_ui(agent):
         progress_bar = gr.Progress()
 
         prompt_template = """
-Analyze the patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown with findings grouped under headings (e.g., 'Missed Diagnoses'). For each finding, include clinical context, risks, and recommendations. Output only markdown bullet points under headings. If no issues, state "No issues identified".
+Analyze the patient record excerpt for clinical oversights. Provide a detailed, evidence-based summary in markdown with findings grouped under headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, Urgent Follow-up. For each finding, include clinical context, risks, and recommendations. Output only markdown bullet points under headings. If no issues, state "No oversights identified" once.
 
-Patient Record Excerpt (Chunk {0} of {1}):
+Patient Record Excerpt:
 {chunk}
 """
 
@@ -302,69 +295,64 @@ Patient Record Excerpt (Chunk {0} of {1}):
                 yield history, None, ""
                 logger.info("Extracted text length: %d chars", len(extracted))
 
-            chunk_size = 3000  # Adjusted for balance
-            chunks = [extracted[i:i + chunk_size] for i in range(0, max(len(extracted), 1), chunk_size)]
-            if not chunks:
-                chunks = [""]  # Ensure at least one chunk
+            chunk_size = 3000
+            chunks = [extracted[i:i + chunk_size] for i in range(0, max(len(extracted), 1), chunk_size)] or [""]
             logger.info("Created %d chunks", len(chunks))
-            combined_response = ""
+            for i, chunk in enumerate(chunks):
+                logger.debug("Chunk %d content: %s...", i + 1, chunk[:100])
+            all_responses = []
             batch_size = 2
 
             try:
                 for batch_idx in range(0, len(chunks), batch_size):
                     batch_chunks = chunks[batch_idx:batch_idx + batch_size]
-                    batch_prompts = [prompt_template.format(i + 1, len(chunks), chunk=chunk[:2000]) for i, chunk in enumerate(batch_chunks)]
+                    batch_prompts = [prompt_template.format(chunk=chunk[:2000]) for chunk in batch_chunks]
                     batch_responses = []
 
                     progress((batch_idx + 1) / len(chunks), desc=f"Analyzing chunks {batch_idx + 1}-{min(batch_idx + batch_size, len(chunks))}/{len(chunks)}")
 
                     async def process_chunk(prompt):
                         chunk_response = ""
+                        raw_outputs = []
                         for chunk_output in agent.run_gradio_chat(
-                            message=prompt, history=[], temperature=0.2, max_new_tokens=256, max_token=1024, call_agent=False, conversation=[]
+                            message=prompt, history=[], temperature=0.2, max_new_tokens=512, max_token=1024, call_agent=False, conversation=[]
                         ):
                             if chunk_output is None:
                                 continue
                             if isinstance(chunk_output, list):
                                 for m in chunk_output:
                                     if hasattr(m, 'content') and m.content:
+                                        raw_outputs.append(m.content)
                                         cleaned = clean_response(m.content)
                                         chunk_response += cleaned + "\n\n"
                             elif isinstance(chunk_output, str) and chunk_output.strip():
+                                raw_outputs.append(chunk_output)
                                 cleaned = clean_response(chunk_output)
                                 chunk_response += cleaned + "\n\n"
+                        logger.debug("Raw outputs: %s", raw_outputs[:100])
                         logger.debug("Chunk response length: %d chars", len(chunk_response))
                         return chunk_response
 
                     futures = [process_chunk(prompt) for prompt in batch_prompts]
                     batch_responses = await asyncio.gather(*futures)
+                    all_responses.extend([resp.strip() for resp in batch_responses if resp.strip()])
                     torch.cuda.empty_cache()
                     gc.collect()
 
-                    for chunk_idx, chunk_response in enumerate(batch_responses, batch_idx + 1):
-                        if chunk_response.strip():
-                            combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"
-                        else:
-                            combined_response += f"--- Analysis for Chunk {chunk_idx} ---\nNo oversights identified for this chunk.\n\n"
-                    history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                    yield history, None, ""
-
-                if combined_response.strip() and not all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")):
-                    history[-1]["content"] = combined_response.strip()
-                else:
-                    history.append({"role": "assistant", "content": "No oversights identified in the provided records."})
+                summary = summarize_findings(all_responses)
+                history.append({"role": "assistant", "content": "Analysis complete. See summary below."})
+                yield history, None, summary
 
-                summary = summarize_findings(combined_response)
                 report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
                 if report_path:
                     with open(report_path, "w", encoding="utf-8") as f:
-                        f.write(combined_response + "\n\n" + summary)
+                        f.write(summary)
                 yield history, report_path if report_path and os.path.exists(report_path) else None, summary
 
             except Exception as e:
                 logger.error("Analysis error: %s", e)
                 history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
-                yield history, None, f"### Summary of Clinical Oversights\nError occurred during analysis: {str(e)}"
+                yield history, None, f"### Comprehensive Clinical Oversight Summary\nError occurred during analysis: {str(e)}"
 
         send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])
         msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])