CPS-Test-Mobile / app.py
Ali2206's picture
Update app.py
bf0551b verified
raw
history blame
9.85 kB
import sys
import os
import pandas as pd
import pdfplumber
import json
import gradio as gr
from typing import List, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed
import hashlib
import shutil
import time
from functools import lru_cache
from threading import Thread
import re
import tempfile
# Environment setup
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)
# Cache directories
base_dir = "/data"
os.makedirs(base_dir, exist_ok=True)
model_cache_dir = os.path.join(base_dir, "txagent_models")
tool_cache_dir = os.path.join(base_dir, "tool_cache")
file_cache_dir = os.path.join(base_dir, "cache")
report_dir = "/data/reports"
vllm_cache_dir = os.path.join(base_dir, "vllm_cache")
os.makedirs(model_cache_dir, exist_ok=True)
os.makedirs(tool_cache_dir, exist_ok=True)
os.makedirs(file_cache_dir, exist_ok=True)
os.makedirs(report_dir, exist_ok=True)
os.makedirs(vllm_cache_dir, exist_ok=True)
os.environ.update({
"TRANSFORMERS_CACHE": model_cache_dir,
"HF_HOME": model_cache_dir,
"VLLM_CACHE_DIR": vllm_cache_dir,
"TOKENIZERS_PARALLELISM": "false",
"CUDA_LAUNCH_BLOCKING": "1"
})
from txagent.txagent import TxAgent
MEDICAL_KEYWORDS = {
'diagnosis', 'assessment', 'plan', 'results', 'medications',
'allergies', 'summary', 'impression', 'findings', 'recommendations'
}
def sanitize_utf8(text: str) -> str:
return text.encode("utf-8", "ignore").decode("utf-8")
def file_hash(path: str) -> str:
with open(path, "rb") as f:
return hashlib.md5(f.read()).hexdigest()
def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
try:
text_chunks = []
with pdfplumber.open(file_path) as pdf:
for i, page in enumerate(pdf.pages[:3]):
text_chunks.append(f"=== Page {i+1} ===\n{(page.extract_text() or '').strip()}")
for i, page in enumerate(pdf.pages[3:max_pages], start=4):
page_text = page.extract_text() or ""
if any(re.search(rf'\\b{kw}\\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
return "\n\n".join(text_chunks)
except Exception as e:
return f"PDF processing error: {str(e)}"
def convert_file_to_json(file_path: str, file_type: str) -> str:
try:
h = file_hash(file_path)
cache_path = os.path.join(file_cache_dir, f"{h}.json")
if os.path.exists(cache_path):
return open(cache_path, "r", encoding="utf-8").read()
if file_type == "pdf":
text = extract_priority_pages(file_path)
result = json.dumps({"filename": os.path.basename(file_path), "content": text, "status": "initial"})
Thread(target=full_pdf_processing, args=(file_path, h)).start()
elif file_type == "csv":
df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str, skip_blank_lines=False, on_bad_lines="skip")
content = df.fillna("").astype(str).values.tolist()
result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
elif file_type in ["xls", "xlsx"]:
try:
df = pd.read_excel(file_path, engine="openpyxl", header=None, dtype=str)
except:
df = pd.read_excel(file_path, engine="xlrd", header=None, dtype=str)
content = df.fillna("").astype(str).values.tolist()
result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
else:
return json.dumps({"error": f"Unsupported file type: {file_type}"})
with open(cache_path, "w", encoding="utf-8") as f:
f.write(result)
return result
except Exception as e:
return json.dumps({"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"})
def full_pdf_processing(file_path: str, file_hash: str):
try:
cache_path = os.path.join(file_cache_dir, f"{file_hash}_full.json")
if os.path.exists(cache_path):
return
with pdfplumber.open(file_path) as pdf:
full_text = "\n".join([f"=== Page {i+1} ===\n{(page.extract_text() or '').strip()}" for i, page in enumerate(pdf.pages)])
result = json.dumps({"filename": os.path.basename(file_path), "content": full_text, "status": "complete"})
with open(cache_path, "w", encoding="utf-8") as f:
f.write(result)
with open(os.path.join(report_dir, f"{file_hash}_report.txt"), "w", encoding="utf-8") as out:
out.write(full_text)
except Exception as e:
print(f"Background processing failed: {str(e)}")
def init_agent():
default_tool_path = os.path.abspath("data/new_tool.json")
target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
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=True,
step_rag_num=8,
seed=100,
additional_default_tools=[],
)
agent.init_model()
return agent
def create_ui(agent: TxAgent):
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
gr.Markdown("<h3 style='text-align: center;'>Identify potential oversights in patient care</h3>")
chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
file_upload = gr.File(label="Upload Medical Records", 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")
conversation_state = gr.State([])
download_output = gr.File(label="Download Full Report")
def analyze_potential_oversights(message: str, history: list, conversation: list, files: list):
start_time = time.time()
try:
history.append({"role": "user", "content": message})
history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
yield history, None
extracted_data = ""
file_hash_value = ""
if files and isinstance(files, list):
with ThreadPoolExecutor(max_workers=4) as executor:
futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files if hasattr(f, 'name')]
extracted_data = "\n".join([sanitize_utf8(f.result()) for f in as_completed(futures)])
file_hash_value = file_hash(files[0].name) if files else ""
analysis_prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
1. List potential missed diagnoses
2. Flag any medication conflicts
3. Note incomplete assessments
4. Highlight abnormal results needing follow-up
Medical Records:\n{extracted_data[:15000]}
### Potential Oversights:\n"""
response = ""
for chunk in agent.run_gradio_chat(
message=analysis_prompt,
history=[],
temperature=0.2,
max_new_tokens=1024,
max_token=4096,
call_agent=False,
conversation=conversation
):
if isinstance(chunk, str):
response += chunk
elif isinstance(chunk, list):
response += "".join([c.content for c in chunk if hasattr(c, 'content')])
history[-1]["content"] = response.replace("[TOOL_CALLS]", "").strip()
yield history, None
final_output = response.replace("[TOOL_CALLS]", "").strip()
if not final_output:
final_output = "No clear oversights identified. Recommend comprehensive review."
report_path = None
if file_hash_value:
possible_report = os.path.join(report_dir, f"{file_hash_value}_report.txt")
if os.path.exists(possible_report):
report_path = possible_report
history[-1] = {"role": "assistant", "content": final_output}
yield history, report_path
except Exception as e:
history.append({"role": "assistant", "content": f"❌ Analysis failed: {str(e)}"})
yield history, None
inputs = [msg_input, chatbot, conversation_state, file_upload]
outputs = [chatbot, download_output]
send_btn.click(analyze_potential_oversights, inputs=inputs, outputs=outputs)
msg_input.submit(analyze_potential_oversights, inputs=inputs, outputs=outputs)
gr.Examples([
["What might have been missed in this patient's treatment?"],
["Are there any medication conflicts in these records?"],
["What abnormal results require follow-up?"]
], inputs=msg_input)
return demo
if __name__ == "__main__":
print("Initializing medical analysis agent...")
agent = init_agent()
print("Launching interface...")
demo = create_ui(agent)
demo.queue().launch(
server_name="0.0.0.0",
server_port=7860,
show_error=True,
allowed_paths=["/data/reports"]
)