Update app.py

app.py

@@ -5,23 +5,37 @@ import pdfplumber
 import json
 import gradio as gr
 from typing import List, Dict, Optional, Generator
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import ProcessPoolExecutor, as_completed
 import hashlib
 import shutil
 import re
 import psutil
-import subprocess
+ مخصوصimport subprocess
 import logging
 import torch
 import gc
 from diskcache import Cache
 import time
 from transformers import AutoTokenizer
+import pyarrow as pa
+import pyarrow.csv as pc
+import pyarrow.parquet as pq
+from vllm import LLM, SamplingParams
+import asyncio
+import threading
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# File handler for response logging
+response_log_file = os.path.join("/data/hf_cache", "response_log.txt")
+response_logger = logging.getLogger("ResponseLogger")
+response_handler = logging.FileHandler(response_log_file, mode="a")
+response_handler.setFormatter(logging.Formatter("%(asctime)s - %(message)s"))
+response_logger.addHandler(response_handler)
+response_logger.setLevel(logging.INFO)
+
 # Persistent directory
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
@@ -61,13 +75,17 @@ def file_hash(path: str) -> str:
         return hashlib.md5(f.read()).hexdigest()
 
 def extract_all_pages(file_path: str, progress_callback=None) -> str:
+    cache_key = f"pdf_{file_hash(file_path)}"
+    if cache_key in cache:
+        return cache[cache_key]
+
     try:
         with pdfplumber.open(file_path) as pdf:
             total_pages = len(pdf.pages)
             if total_pages == 0:
                 return ""
 
-        batch_size = 10
+        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
@@ -77,11 +95,11 @@ def extract_all_pages(file_path: str, progress_callback=None) -> str:
             with pdfplumber.open(file_path) as pdf:
                 for page in pdf.pages[start:end]:
                     page_num = start + pdf.pages.index(page)
-                    page_text = page.extract_text() or ""
+                    page_text = page.extract_text_simple() or ""
                     results.append((page_num, f"=== Page {page_num + 1} ===\n{page_text.strip()}"))
             return results
 
-        with ThreadPoolExecutor(max_workers=6) as executor:
+        with ProcessPoolExecutor(max_workers=4) as executor:
             futures = [executor.submit(extract_batch, start, end) for start, end in batches]
             for future in as_completed(futures):
                 for page_num, text in future.result():
@@ -90,62 +108,54 @@ def extract_all_pages(file_path: str, progress_callback=None) -> str:
                 if progress_callback:
                     progress_callback(min(processed_pages, total_pages), total_pages)
 
-        return "\n\n".join(filter(None, text_chunks))
+        result = "\n\n".join(filter(None, text_chunks))
+        cache[cache_key] = result
+        return result
     except Exception as e:
         logger.error("PDF processing error: %s", e)
         return f"PDF processing error: {str(e)}"
 
 def excel_to_json(file_path: str) -> List[Dict]:
-    """Convert Excel file to JSON with optimized processing"""
+    cache_key = f"excel_{file_hash(file_path)}"
+    if cache_key in cache:
+        return cache[cache_key]
+
     try:
-        # First try with openpyxl (faster for xlsx)
-        try:
-            df = pd.read_excel(file_path, engine='openpyxl', header=None, dtype=str)
-        except Exception:
-            # Fall back to xlrd if needed
-            df = pd.read_excel(file_path, engine='xlrd', header=None, dtype=str)
-        
-        # Convert to list of lists with null handling
+        table = pq.read_table(file_path)
+        df = table.to_pandas(use_threads=True, split_blocks=True)
         content = df.where(pd.notnull(df), "").astype(str).values.tolist()
-        
-        return [{
+        result = [{
             "filename": os.path.basename(file_path),
             "rows": content,
             "type": "excel"
         }]
+        cache[cache_key] = result
+        return result
     except Exception as e:
         logger.error(f"Error processing Excel file: {e}")
         return [{"error": f"Error processing Excel file: {str(e)}"}]
 
 def csv_to_json(file_path: str) -> List[Dict]:
-    """Convert CSV file to JSON with optimized processing"""
+    cache_key = f"csv_{file_hash(file_path)}"
+    if cache_key in cache:
+        return cache[cache_key]
+
     try:
-        # Read CSV in chunks if large
-        chunks = []
-        for chunk in pd.read_csv(
-            file_path,
-            header=None,
-            dtype=str,
-            encoding_errors='replace',
-            on_bad_lines='skip',
-            chunksize=10000
-        ):
-            chunks.append(chunk)
-        
-        df = pd.concat(chunks) if chunks else pd.DataFrame()
+        table = pc.read_csv(file_path, parse_options=pc.ParseOptions(invalid_row_handler=lambda x: "skip"))
+        df = table.to_pandas(use_threads=True, split_blocks=True)
         content = df.where(pd.notnull(df), "").astype(str).values.tolist()
-        
-        return [{
+        result = [{
             "filename": os.path.basename(file_path),
             "rows": content,
             "type": "csv"
         }]
+        cache[cache_key] = result
+        return result
     except Exception as e:
         logger.error(f"Error processing CSV file: {e}")
         return [{"error": f"Error processing CSV file: {str(e)}"}]
 
 def process_file(file_path: str, file_type: str) -> List[Dict]:
-    """Process file based on type and return JSON data"""
     try:
         if file_type == "pdf":
             text = extract_all_pages(file_path)
@@ -165,18 +175,22 @@ def process_file(file_path: str, file_type: str) -> List[Dict]:
         logger.error("Error processing %s: %s", os.path.basename(file_path), e)
         return [{"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"}]
 
-def tokenize_and_chunk(text: str, max_tokens: int = 1800) -> List[str]:
-    """Split text into chunks based on token count"""
-    tokens = tokenizer.encode(text)
+def tokenize_and_chunk(text: str, max_tokens: int = 800) -> List[str]:
+    cache_key = f"tokens_{hashlib.md5(text.encode()).hexdigest()}"
+    if cache_key in cache:
+        return cache[cache_key]
+
+    tokens = tokenizer.encode(text, add_special_tokens=False)
     chunks = []
     for i in range(0, len(tokens), max_tokens):
         chunk_tokens = tokens[i:i + max_tokens]
-        chunks.append(tokenizer.decode(chunk_tokens))
+        chunks.append(tokenizer.decode(chunk_tokens, skip_special_tokens=True))
+    cache[cache_key] = chunks
     return chunks
 
 def log_system_usage(tag=""):
     try:
-        cpu = psutil.cpu_percent(interval=1)
+        cpu = psutil.cpu_percent(interval=0.1)
         mem = psutil.virtual_memory()
         logger.info("[%s] CPU: %.1f%% | RAM: %dMB / %dMB", tag, cpu, mem.used // (1024**2), mem.total // (1024**2))
         result = subprocess.run(
@@ -261,27 +275,27 @@ def init_agent():
     if not os.path.exists(target_tool_path):
         shutil.copy(default_tool_path, target_tool_path)
 
-    agent = TxAgent(
-        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
-        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-        tool_files_dict={"new_tool": target_tool_path},
-        force_finish=True,
-        enable_checker=False,
-        step_rag_num=4,
-        seed=100,
-        additional_default_tools=[],
+    llm = LLM(
+        model="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+        gpu_memory_utilization=0.8,
+        max_model_len=2048,
+        tensor_parallel_size=1,
+    )
+    sampling_params = SamplingParams(
+        temperature=0.2,
+        max_tokens=256,  # Reduced for faster streaming
+        stop=["</s>", "[INST]"],
     )
-    agent.init_model()
     log_system_usage("After Load")
     logger.info("Agent Ready")
-    return agent
+    return llm, sampling_params
 
-def create_ui(agent):
+async def create_ui(llm, sampling_params):
     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 Missed Diagnoses")
-        file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
+        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")
         download_output = gr.File(label="Download Full Report")
@@ -293,7 +307,10 @@ Patient Record Excerpt (Chunk {0} of {1}):
 {chunk}
 """
 
-        def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
+        def log_response_partial(text: str):
+            response_logger.info(text)
+
+        async def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
             history.append({"role": "user", "content": message})
             yield history, None, ""
 
@@ -301,8 +318,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
             file_hash_value = ""
             
             if files:
-                # Process files in parallel
-                with ThreadPoolExecutor(max_workers=4) as executor:
+                with ProcessPoolExecutor(max_workers=4) as executor:
                     futures = []
                     for f in files:
                         file_type = f.name.split(".")[-1].lower()
@@ -323,14 +339,11 @@ Patient Record Excerpt (Chunk {0} of {1}):
                 history.append({"role": "assistant", "content": "✅ File processing complete"})
                 yield history, None, ""
 
-            # Convert extracted data to JSON text
             text_content = "\n".join(json.dumps(item) for item in extracted)
-            
-            # Tokenize and chunk the content properly
             chunks = tokenize_and_chunk(text_content)
             combined_response = ""
-            batch_size = 2  # Reduced batch size to prevent token overflow
-            
+            batch_size = 1
+
             try:
                 for batch_idx in range(0, len(chunks), batch_size):
                     batch_chunks = chunks[batch_idx:batch_idx + batch_size]
@@ -338,7 +351,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                         prompt_template.format(
                             batch_idx + i + 1,
                             len(chunks),
-                            chunk=chunk[:1800]  # Conservative chunk size
+                            chunk=chunk[:800]
                         )
                         for i, chunk in enumerate(batch_chunks)
                     ]
@@ -346,63 +359,54 @@ Patient Record Excerpt (Chunk {0} of {1}):
                     progress((batch_idx) / len(chunks), 
                            desc=f"Analyzing batch {(batch_idx // batch_size) + 1}/{(len(chunks) + batch_size - 1) // batch_size}")
                     
-                    # Process batch in parallel
-                    with ThreadPoolExecutor(max_workers=len(batch_prompts)) as executor:
-                        future_to_prompt = {
-                            executor.submit(
-                                agent.run_gradio_chat,
-                                prompt, [], 0.2, 512, 2048, False, []
-                            ): prompt
-                            for prompt in batch_prompts
-                        }
-                        
-                        for future in as_completed(future_to_prompt):
+                    with torch.no_grad():
+                        for prompt in batch_prompts:
                             chunk_response = ""
-                            for chunk_output in future.result():
-                                if chunk_output is None:
-                                    continue
-                                if isinstance(chunk_output, list):
-                                    for m in chunk_output:
-                                        if hasattr(m, 'content') and m.content:
-                                            cleaned = clean_response(m.content)
-                                            if cleaned:
-                                                chunk_response += cleaned + " "
-                                elif isinstance(chunk_output, str) and chunk_output.strip():
-                                    cleaned = clean_response(chunk_output)
-                                    if cleaned:
-                                        chunk_response += cleaned + " "
-                            
-                            combined_response += f"--- Analysis for Chunk {batch_idx + 1} ---\n{chunk_response.strip()}\n"
-                            history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                            yield history, None, ""
-                            
-                            # Clean up memory
-                            torch.cuda.empty_cache()
-                            gc.collect()
-
-                # Generate final summary
+                            current_response = ""
+                            stream = llm.generate([prompt], sampling_params, use_tqdm=False)
+                            for output in stream:
+                                for request_output in output:
+                                    new_text = request_output.outputs[0].text[len(current_response):]
+                                    if new_text:
+                                        current_response += new_text
+                                        cleaned = clean_response(current_response)
+                                        if cleaned and cleaned != chunk_response:
+                                            chunk_response = cleaned
+                                            history[-1] = {"role": "assistant", "content": chunk_response}
+                                            threading.Thread(target=log_response_partial, args=(chunk_response,)).start()
+                                            yield history, None, ""
+                                            await asyncio.sleep(0.01)  # Prevent UI blocking
+                    
+                            if chunk_response:
+                                combined_response += f"--- Analysis for Chunk {batch_idx + 1} ---\n{chunk_response}\n"
+                    
+                    torch.cuda.empty_cache()
+                    gc.collect()
+
                 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)
+                    threading.Thread(target=log_response_partial, args=(summary,)).start()
                 
                 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)}"})
+                threading.Thread(target=log_response_partial, args=(f"Error: {str(e)}",)).start()
                 yield history, None, f"Error 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])
+        send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary], _js="() => {return {streaming: true}}")
+        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary], _js="() => {return {streaming: true}}")
     return demo
 
 if __name__ == "__main__":
     try:
         logger.info("Launching app...")
-        agent = init_agent()
-        demo = create_ui(agent)
+        llm, sampling_params = init_agent()
+        demo = asyncio.run(create_ui(llm, sampling_params))
         demo.queue(api_open=False).launch(
             server_name="0.0.0.0",
             server_port=7860,