|
|
import sys |
|
|
import os |
|
|
import pandas as pd |
|
|
import json |
|
|
import gradio as gr |
|
|
from typing import List, Tuple, Dict, Any, Union |
|
|
import hashlib |
|
|
import shutil |
|
|
import re |
|
|
from datetime import datetime |
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed |
|
|
|
|
|
|
|
|
persistent_dir = "/data/hf_cache" |
|
|
os.makedirs(persistent_dir, exist_ok=True) |
|
|
|
|
|
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 |
|
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src"))) |
|
|
from txagent.txagent import TxAgent |
|
|
|
|
|
MAX_MODEL_TOKENS = 32768 |
|
|
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] |
|
|
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]: |
|
|
effective_max = max_tokens - PROMPT_OVERHEAD |
|
|
lines, chunks, curr_chunk, curr_tokens = text.split("\n"), [], [], 0 |
|
|
for line in lines: |
|
|
t = estimate_tokens(line) |
|
|
if curr_tokens + t > effective_max: |
|
|
if curr_chunk: |
|
|
chunks.append("\n".join(curr_chunk)) |
|
|
curr_chunk, curr_tokens = [line], t |
|
|
else: |
|
|
curr_chunk.append(line) |
|
|
curr_tokens += t |
|
|
if curr_chunk: |
|
|
chunks.append("\n".join(curr_chunk)) |
|
|
return chunks |
|
|
|
|
|
def build_prompt_from_text(chunk: str) -> str: |
|
|
return f""" |
|
|
### Unstructured Clinical Records |
|
|
|
|
|
Analyze the following clinical notes and provide a detailed, concise summary focusing on: |
|
|
- Diagnostic Patterns |
|
|
- Medication Issues |
|
|
- Missed Opportunities |
|
|
- Inconsistencies |
|
|
- Follow-up Recommendations |
|
|
|
|
|
--- |
|
|
|
|
|
{chunk} |
|
|
|
|
|
--- |
|
|
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): |
|
|
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 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"): |
|
|
return messages + [{"role": "assistant", "content": "❌ Please upload a valid Excel file."}], None |
|
|
|
|
|
messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"}) |
|
|
text = extract_text_from_excel(file.name) |
|
|
chunks = split_text_into_chunks(text) |
|
|
chunk_responses = [None] * len(chunks) |
|
|
|
|
|
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=[]): |
|
|
if isinstance(res, str): |
|
|
response += res |
|
|
elif hasattr(res, "content"): |
|
|
response += res.content |
|
|
elif isinstance(res, list): |
|
|
for r in res: |
|
|
if hasattr(r, "content"): |
|
|
response += r.content |
|
|
return i, clean_response(response) |
|
|
|
|
|
with ThreadPoolExecutor(max_workers=1) 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() |
|
|
chunk_responses[i] = result |
|
|
|
|
|
valid = [r for r in chunk_responses if r and not r.startswith("❌")] |
|
|
if not valid: |
|
|
return messages + [{"role": "assistant", "content": "❌ No valid chunk results."}], None |
|
|
|
|
|
summary_prompt = f"Summarize this analysis in a final structured report:\n\n" + "\n\n".join(valid) |
|
|
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=[]): |
|
|
if isinstance(res, str): |
|
|
final_report += res |
|
|
elif hasattr(res, "content"): |
|
|
final_report += res.content |
|
|
|
|
|
cleaned = clean_response(final_report) |
|
|
report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md") |
|
|
with open(report_path, 'w') as f: |
|
|
f.write(f"# 🧠 Final Patient Report\n\n{cleaned}") |
|
|
|
|
|
messages.append({"role": "assistant", "content": f"📊 Final Report:\n\n{cleaned}"}) |
|
|
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=""" |
|
|
:root { |
|
|
--primary-color: #2563eb; |
|
|
--secondary-color: #1e40af; |
|
|
--bg-color: #f8fafc; |
|
|
--text-color: #1e293b; |
|
|
--user-bubble: #ffffff; |
|
|
--bot-bubble: #f1f5f9; |
|
|
--border-color: #e2e8f0; |
|
|
} |
|
|
body { |
|
|
font-family: 'Inter', sans-serif; |
|
|
background-color: var(--bg-color); |
|
|
color: var(--text-color); |
|
|
} |
|
|
.gradio-container { |
|
|
max-width: 800px; |
|
|
margin: 0 auto; |
|
|
padding: 20px; |
|
|
background-color: var(--bg-color); |
|
|
} |
|
|
.chat-container { |
|
|
display: flex; |
|
|
flex-direction: column; |
|
|
height: 80vh; |
|
|
border: 1px solid var(--border-color); |
|
|
border-radius: 12px; |
|
|
overflow: hidden; |
|
|
background-color: white; |
|
|
} |
|
|
.chat-header { |
|
|
padding: 16px; |
|
|
background-color: white; |
|
|
border-bottom: 1px solid var(--border-color); |
|
|
font-weight: 600; |
|
|
font-size: 18px; |
|
|
} |
|
|
.chat-messages { |
|
|
flex: 1; |
|
|
padding: 16px; |
|
|
overflow-y: auto; |
|
|
background-color: var(--bg-color); |
|
|
} |
|
|
.message { |
|
|
display: flex; |
|
|
margin-bottom: 16px; |
|
|
align-items: flex-start; |
|
|
} |
|
|
.message-avatar { |
|
|
width: 36px; |
|
|
height: 36px; |
|
|
border-radius: 50%; |
|
|
margin-right: 12px; |
|
|
background-color: var(--primary-color); |
|
|
color: white; |
|
|
display: flex; |
|
|
align-items: center; |
|
|
justify-content: center; |
|
|
font-weight: bold; |
|
|
} |
|
|
.message-content { |
|
|
max-width: 80%; |
|
|
} |
|
|
.user-message .message-content { |
|
|
margin-left: auto; |
|
|
background-color: var(--user-bubble); |
|
|
padding: 12px 16px; |
|
|
border-radius: 18px 18px 0 18px; |
|
|
box-shadow: 0 1px 2px rgba(0,0,0,0.1); |
|
|
} |
|
|
.bot-message .message-content { |
|
|
background-color: var(--bot-bubble); |
|
|
padding: 12px 16px; |
|
|
border-radius: 18px 18px 18px 0; |
|
|
box-shadow: 0 1px 2px rgba(0,0,0,0.1); |
|
|
} |
|
|
.message-time { |
|
|
font-size: 12px; |
|
|
color: #64748b; |
|
|
margin-top: 4px; |
|
|
} |
|
|
.chat-input-container { |
|
|
padding: 16px; |
|
|
background-color: white; |
|
|
border-top: 1px solid var(--border-color); |
|
|
} |
|
|
.file-upload { |
|
|
display: flex; |
|
|
gap: 8px; |
|
|
margin-bottom: 12px; |
|
|
} |
|
|
.upload-btn { |
|
|
flex: 1; |
|
|
} |
|
|
.analyze-btn { |
|
|
background-color: var(--primary-color); |
|
|
color: white; |
|
|
border: none; |
|
|
padding: 10px 16px; |
|
|
border-radius: 8px; |
|
|
cursor: pointer; |
|
|
font-weight: 500; |
|
|
} |
|
|
.analyze-btn:hover { |
|
|
background-color: var(--secondary-color); |
|
|
} |
|
|
.report-download { |
|
|
margin-top: 12px; |
|
|
padding: 12px; |
|
|
background-color: var(--bot-bubble); |
|
|
border-radius: 8px; |
|
|
display: none; |
|
|
} |
|
|
""") as demo: |
|
|
with gr.Column(elem_classes="chat-container"): |
|
|
gr.Markdown("""<div class="chat-header">Patient History AI Assistant</div>""") |
|
|
|
|
|
chatbot = gr.Chatbot( |
|
|
elem_classes="chat-messages", |
|
|
label=None, |
|
|
show_label=False, |
|
|
bubble_full_width=False, |
|
|
avatar_images=[ |
|
|
"https://ui-avatars.com/api/?name=AI&background=2563eb&color=fff&size=128", |
|
|
None |
|
|
] |
|
|
) |
|
|
|
|
|
with gr.Column(elem_classes="chat-input-container"): |
|
|
with gr.Row(elem_classes="file-upload"): |
|
|
file_upload = gr.File( |
|
|
label="Upload Excel File", |
|
|
file_types=[".xlsx"], |
|
|
elem_classes="upload-btn", |
|
|
scale=4 |
|
|
) |
|
|
analyze_btn = gr.Button( |
|
|
"Analyze", |
|
|
elem_classes="analyze-btn", |
|
|
scale=1 |
|
|
) |
|
|
|
|
|
report_output = gr.File( |
|
|
label="Download Report", |
|
|
visible=False, |
|
|
interactive=False, |
|
|
elem_classes="report-download" |
|
|
) |
|
|
|
|
|
chatbot_state = gr.State(value=[]) |
|
|
|
|
|
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 |
|
|
|
|
|
analyze_btn.click( |
|
|
fn=update_ui, |
|
|
inputs=[file_upload, chatbot_state], |
|
|
outputs=[chatbot, report_output, chatbot_state] |
|
|
) |
|
|
|
|
|
return demo |
|
|
|
|
|
if __name__ == "__main__": |
|
|
try: |
|
|
agent = init_agent() |
|
|
demo = create_ui(agent) |
|
|
demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False) |
|
|
except Exception as e: |
|
|
print(f"Error: {str(e)}") |
|
|
sys.exit(1) |
