|
|
import sys |
|
|
import os |
|
|
import json |
|
|
import shutil |
|
|
import re |
|
|
import gc |
|
|
import time |
|
|
from datetime import datetime |
|
|
from typing import List, Tuple, Dict, Union |
|
|
import pandas as pd |
|
|
import pdfplumber |
|
|
import gradio as gr |
|
|
import torch |
|
|
import matplotlib.pyplot as plt |
|
|
from fpdf import FPDF |
|
|
import unicodedata |
|
|
|
|
|
|
|
|
persistent_dir = "/data/hf_cache" |
|
|
model_cache_dir = os.path.join(persistent_dir, "txagent_models") |
|
|
tool_cache_dir = os.path.join(persistent_dir, "tool_cache") |
|
|
file_cache_dir = os.path.join(persistent_dir, "cache") |
|
|
report_dir = os.path.join(persistent_dir, "reports") |
|
|
|
|
|
for d in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]: |
|
|
os.makedirs(d, exist_ok=True) |
|
|
|
|
|
os.environ["HF_HOME"] = model_cache_dir |
|
|
os.environ["TRANSFORMERS_CACHE"] = model_cache_dir |
|
|
|
|
|
current_dir = os.path.dirname(os.path.abspath(__file__)) |
|
|
src_path = os.path.abspath(os.path.join(current_dir, "src")) |
|
|
sys.path.insert(0, src_path) |
|
|
|
|
|
from txagent.txagent import TxAgent |
|
|
|
|
|
MAX_MODEL_TOKENS = 131072 |
|
|
MAX_NEW_TOKENS = 4096 |
|
|
MAX_CHUNK_TOKENS = 8192 |
|
|
BATCH_SIZE = 2 |
|
|
PROMPT_OVERHEAD = 300 |
|
|
SAFE_SLEEP = 0.5 |
|
|
|
|
|
def estimate_tokens(text: str) -> int: |
|
|
return len(text) // 4 + 1 |
|
|
|
|
|
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) |
|
|
return text.strip() |
|
|
|
|
|
def remove_duplicate_paragraphs(text: str) -> str: |
|
|
paragraphs = text.strip().split("\n\n") |
|
|
seen = set() |
|
|
unique_paragraphs = [] |
|
|
for p in paragraphs: |
|
|
clean_p = p.strip() |
|
|
if clean_p and clean_p not in seen: |
|
|
unique_paragraphs.append(clean_p) |
|
|
seen.add(clean_p) |
|
|
return "\n\n".join(unique_paragraphs) |
|
|
|
|
|
def extract_text_from_excel(path: str) -> str: |
|
|
all_text = [] |
|
|
xls = pd.ExcelFile(path) |
|
|
for sheet_name in xls.sheet_names: |
|
|
try: |
|
|
df = xls.parse(sheet_name).astype(str).fillna("") |
|
|
except Exception: |
|
|
continue |
|
|
for _, row in df.iterrows(): |
|
|
non_empty = [cell.strip() for cell in row if cell.strip()] |
|
|
if len(non_empty) >= 2: |
|
|
text_line = " | ".join(non_empty) |
|
|
if len(text_line) > 15: |
|
|
all_text.append(f"[{sheet_name}] {text_line}") |
|
|
return "\n".join(all_text) |
|
|
|
|
|
def extract_text_from_csv(path: str) -> str: |
|
|
all_text = [] |
|
|
try: |
|
|
df = pd.read_csv(path).astype(str).fillna("") |
|
|
except Exception: |
|
|
return "" |
|
|
for _, row in df.iterrows(): |
|
|
non_empty = [cell.strip() for cell in row if cell.strip()] |
|
|
if len(non_empty) >= 2: |
|
|
text_line = " | ".join(non_empty) |
|
|
if len(text_line) > 15: |
|
|
all_text.append(text_line) |
|
|
return "\n".join(all_text) |
|
|
|
|
|
def extract_text_from_pdf(path: str) -> str: |
|
|
import logging |
|
|
logging.getLogger("pdfminer").setLevel(logging.ERROR) |
|
|
all_text = [] |
|
|
try: |
|
|
with pdfplumber.open(path) as pdf: |
|
|
for page in pdf.pages: |
|
|
text = page.extract_text() |
|
|
if text: |
|
|
all_text.append(text.strip()) |
|
|
except Exception: |
|
|
return "" |
|
|
return "\n".join(all_text) |
|
|
|
|
|
def extract_text(file_path: str) -> str: |
|
|
if file_path.endswith(".xlsx"): |
|
|
return extract_text_from_excel(file_path) |
|
|
elif file_path.endswith(".csv"): |
|
|
return extract_text_from_csv(file_path) |
|
|
elif file_path.endswith(".pdf"): |
|
|
return extract_text_from_pdf(file_path) |
|
|
else: |
|
|
return "" |
|
|
|
|
|
def split_text(text: str, max_tokens=MAX_CHUNK_TOKENS) -> List[str]: |
|
|
effective_limit = max_tokens - PROMPT_OVERHEAD |
|
|
chunks, current, current_tokens = [], [], 0 |
|
|
for line in text.split("\n"): |
|
|
tokens = estimate_tokens(line) |
|
|
if current_tokens + tokens > effective_limit: |
|
|
if current: |
|
|
chunks.append("\n".join(current)) |
|
|
current, current_tokens = [line], tokens |
|
|
else: |
|
|
current.append(line) |
|
|
current_tokens += tokens |
|
|
if current: |
|
|
chunks.append("\n".join(current)) |
|
|
return chunks |
|
|
|
|
|
def batch_chunks(chunks: List[str], batch_size: int = BATCH_SIZE) -> List[List[str]]: |
|
|
return [chunks[i:i+batch_size] for i in range(0, len(chunks), batch_size)] |
|
|
|
|
|
def build_prompt(chunk: str) -> str: |
|
|
return f"""### Unstructured Clinical Records\n\nAnalyze the clinical notes below and summarize with:\n- Diagnostic Patterns\n- Medication Issues\n- Missed Opportunities\n- Inconsistencies\n- Follow-up Recommendations\n\n---\n\n{chunk}\n\n---\nRespond concisely in bullet points with clinical reasoning.""" |
|
|
|
|
|
def init_agent() -> TxAgent: |
|
|
tool_path = os.path.join(tool_cache_dir, "new_tool.json") |
|
|
if not os.path.exists(tool_path): |
|
|
shutil.copy(os.path.abspath("data/new_tool.json"), 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": tool_path}, |
|
|
force_finish=True, |
|
|
enable_checker=True, |
|
|
step_rag_num=4, |
|
|
seed=100 |
|
|
) |
|
|
agent.init_model() |
|
|
return agent |
|
|
|
|
|
def analyze_batches(agent, batches: List[List[str]]) -> List[str]: |
|
|
results = [] |
|
|
for batch in batches: |
|
|
prompt = "\n\n".join(build_prompt(chunk) for chunk in batch) |
|
|
try: |
|
|
batch_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, |
|
|
call_agent=False, |
|
|
conversation=[] |
|
|
): |
|
|
if isinstance(r, str): |
|
|
batch_response += r |
|
|
elif isinstance(r, list): |
|
|
for m in r: |
|
|
if hasattr(m, "content"): |
|
|
batch_response += m.content |
|
|
elif hasattr(r, "content"): |
|
|
batch_response += r.content |
|
|
results.append(clean_response(batch_response)) |
|
|
time.sleep(SAFE_SLEEP) |
|
|
except Exception as e: |
|
|
results.append(f"❌ Batch failed: {str(e)}") |
|
|
time.sleep(SAFE_SLEEP * 2) |
|
|
torch.cuda.empty_cache() |
|
|
gc.collect() |
|
|
return results |
|
|
|
|
|
def generate_final_summary(agent, combined: str) -> str: |
|
|
combined = remove_duplicate_paragraphs(combined) |
|
|
final_prompt = f""" |
|
|
You are an expert clinical summarizer. Analyze the following summaries carefully and generate a **single final concise structured medical report**, avoiding any repetition or redundancy. |
|
|
|
|
|
Summaries: |
|
|
{combined} |
|
|
|
|
|
Respond with: |
|
|
- Diagnostic Patterns |
|
|
- Medication Issues |
|
|
- Missed Opportunities |
|
|
- Inconsistencies |
|
|
- Follow-up Recommendations |
|
|
Avoid repeating the same points multiple times. |
|
|
""".strip() |
|
|
|
|
|
final_response = "" |
|
|
for r in agent.run_gradio_chat( |
|
|
message=final_prompt, |
|
|
history=[], |
|
|
temperature=0.0, |
|
|
max_new_tokens=MAX_NEW_TOKENS, |
|
|
max_token=MAX_MODEL_TOKENS, |
|
|
call_agent=False, |
|
|
conversation=[] |
|
|
): |
|
|
if isinstance(r, str): |
|
|
final_response += r |
|
|
elif isinstance(r, list): |
|
|
for m in r: |
|
|
if hasattr(m, "content"): |
|
|
final_response += m.content |
|
|
elif hasattr(r, "content"): |
|
|
final_response += r.content |
|
|
|
|
|
final_response = clean_response(final_response) |
|
|
final_response = remove_duplicate_paragraphs(final_response) |
|
|
return final_response |
|
|
|
|
|
def remove_non_ascii(text): |
|
|
return unicodedata.normalize('NFKD', text).encode('ascii', 'ignore').decode('ascii') |
|
|
|
|
|
def generate_pdf_report_with_charts(summary: str, report_path: str): |
|
|
chart_dir = os.path.join(os.path.dirname(report_path), "charts") |
|
|
os.makedirs(chart_dir, exist_ok=True) |
|
|
|
|
|
chart_path = os.path.join(chart_dir, "summary_chart.png") |
|
|
categories = ['Diagnostics', 'Medications', 'Missed', 'Inconsistencies', 'Follow-up'] |
|
|
values = [4, 2, 3, 1, 5] |
|
|
plt.figure(figsize=(6, 4)) |
|
|
plt.bar(categories, values) |
|
|
plt.title('Clinical Issues Overview') |
|
|
plt.tight_layout() |
|
|
plt.savefig(chart_path) |
|
|
plt.close() |
|
|
|
|
|
pdf_path = report_path.replace('.md', '.pdf') |
|
|
pdf = FPDF() |
|
|
pdf.add_page() |
|
|
pdf.set_font("Arial", size=12) |
|
|
pdf.multi_cell(0, 10, txt="Final Medical Report", align="C") |
|
|
pdf.ln(5) |
|
|
|
|
|
for line in summary.split("\n"): |
|
|
clean_line = remove_non_ascii(line) |
|
|
pdf.multi_cell(0, 10, txt=clean_line) |
|
|
|
|
|
pdf.ln(10) |
|
|
pdf.image(chart_path, w=150) |
|
|
pdf.output(pdf_path) |
|
|
return pdf_path |
|
|
|
|
|
def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]: |
|
|
if not file or not hasattr(file, "name"): |
|
|
messages.append({"role": "assistant", "content": "❌ Please upload a valid file."}) |
|
|
return messages, None |
|
|
|
|
|
start_time = time.time() |
|
|
messages.append({"role": "user", "content": f"📂 Processing file: {os.path.basename(file.name)}"}) |
|
|
|
|
|
try: |
|
|
extracted = extract_text(file.name) |
|
|
if not extracted: |
|
|
messages.append({"role": "assistant", "content": "❌ Could not extract text."}) |
|
|
return messages, None |
|
|
|
|
|
chunks = split_text(extracted) |
|
|
batches = batch_chunks(chunks, batch_size=BATCH_SIZE) |
|
|
messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Analyzing..."}) |
|
|
|
|
|
batch_results = analyze_batches(agent, batches) |
|
|
valid = [res for res in batch_results if not res.startswith("❌")] |
|
|
|
|
|
if not valid: |
|
|
messages.append({"role": "assistant", "content": "❌ No valid batch outputs."}) |
|
|
return messages, None |
|
|
|
|
|
summary = generate_final_summary(agent, "\n\n".join(valid)) |
|
|
|
|
|
report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md") |
|
|
with open(report_path, 'w', encoding='utf-8') as f: |
|
|
f.write(f"# Final Medical Report\n\n{summary}") |
|
|
|
|
|
pdf_path = generate_pdf_report_with_charts(summary, report_path) |
|
|
|
|
|
end_time = time.time() |
|
|
elapsed_time = end_time - start_time |
|
|
|
|
|
messages.append({"role": "assistant", "content": f"📊 **Final Report:**\n\n{summary}"}) |
|
|
messages.append({"role": "assistant", "content": f"✅ Report generated in **{elapsed_time:.2f} seconds**.\n\n📥 PDF report ready: {os.path.basename(pdf_path)}"}) |
|
|
|
|
|
return messages, pdf_path |
|
|
|
|
|
except Exception as e: |
|
|
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: #0e1621; color: #e0e0e0; padding: 16px; } |
|
|
button.svelte-1ipelgc { background: linear-gradient(to right, #1e88e5, #0d47a1) !important; border: 1px solid #0d47a1 !important; color: white !important; font-weight: bold !important; padding: 10px 20px !important; border-radius: 8px !important; } |
|
|
button.svelte-1ipelgc:hover { background: linear-gradient(to right, #2196f3, #1565c0) !important; border: 1px solid #1565c0 !important; } |
|
|
.gr-column { align-items: center !important; gap: 12px; } |
|
|
.gr-file, .gr-button { width: 100% !important; max-width: 400px; } |
|
|
""") as demo: |
|
|
gr.Markdown(""" |
|
|
<h2 style='text-align:center;'>📄 CPS: Clinical Patient Support System</h2> |
|
|
<p style='text-align:center;'>Analyze and summarize unstructured medical files using AI (optimized for A100 GPU).</p> |
|
|
""") |
|
|
|
|
|
with gr.Column(): |
|
|
chatbot = gr.Chatbot(label="🧠 CPS Assistant", height=480, type="messages") |
|
|
upload = gr.File(label="📂 Upload Medical File", file_types=[".xlsx", ".csv", ".pdf"]) |
|
|
analyze = gr.Button("🧠 Analyze") |
|
|
download = gr.File(label="📥 Download Report", visible=False, interactive=False) |
|
|
|
|
|
state = gr.State(value=[]) |
|
|
|
|
|
def handle_analysis(file, chat): |
|
|
messages, report_path = process_report(agent, file, chat) |
|
|
return messages, gr.update(visible=bool(report_path), value=report_path), messages |
|
|
|
|
|
analyze.click(fn=handle_analysis, inputs=[upload, state], outputs=[chatbot, download, state]) |
|
|
|
|
|
return demo |
|
|
|
|
|
if __name__ == "__main__": |
|
|
agent = init_agent() |
|
|
ui = create_ui(agent) |
|
|
ui.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False) |
