|
|
import sys |
|
|
import os |
|
|
import pandas as pd |
|
|
import gradio as gr |
|
|
from typing import List, Tuple |
|
|
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") |
|
|
report_dir = os.path.join(persistent_dir, "reports") |
|
|
|
|
|
for d in [model_cache_dir, tool_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 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([cell for cell in row if cell.strip()]), axis=1) |
|
|
sheet_text = [f"[{sheet_name}] {line}" for line in rows if line.strip()] |
|
|
all_text.extend(sheet_text) |
|
|
return "\n".join(all_text) |
|
|
|
|
|
def split_text_into_chunks(text: str) -> List[str]: |
|
|
effective_max = MAX_CHUNK_TOKENS - PROMPT_OVERHEAD |
|
|
lines, chunks, curr_chunk = text.split("\n"), [], [] |
|
|
curr_tokens = sum(len(line.split()) for line in curr_chunk) |
|
|
|
|
|
for line in lines: |
|
|
line_tokens = len(line.split()) |
|
|
if curr_tokens + line_tokens > effective_max: |
|
|
if curr_chunk: |
|
|
chunks.append("\n".join(curr_chunk)) |
|
|
curr_chunk, curr_tokens = [line], line_tokens |
|
|
else: |
|
|
curr_chunk.append(line) |
|
|
curr_tokens += line_tokens |
|
|
if curr_chunk: |
|
|
chunks.append("\n".join(curr_chunk)) |
|
|
return chunks |
|
|
|
|
|
def build_prompt_from_text(chunk: str) -> str: |
|
|
return f"""Analyze these clinical notes and provide: |
|
|
- Diagnostic patterns |
|
|
- Medication issues |
|
|
- Missed opportunities |
|
|
- Inconsistencies |
|
|
- Follow-up recommendations |
|
|
|
|
|
Respond with clear bullet points: |
|
|
|
|
|
{chunk}""" |
|
|
|
|
|
def init_agent(): |
|
|
tool_path = os.path.join(tool_cache_dir, "new_tool.json") |
|
|
if not os.path.exists(tool_path): |
|
|
import shutil |
|
|
shutil.copy("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[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], str]: |
|
|
messages = chatbot_state.copy() if chatbot_state else [] |
|
|
|
|
|
if file is None: |
|
|
messages.append(("assistant", "❌ Please upload a valid Excel file.")) |
|
|
return messages, None |
|
|
|
|
|
messages.append(("user", f"Processing Excel file: {os.path.basename(file.name)}")) |
|
|
yield messages, None |
|
|
|
|
|
try: |
|
|
text = extract_text_from_excel(file.name) |
|
|
chunks = split_text_into_chunks(text) |
|
|
|
|
|
messages.append(("assistant", "🔍 Analyzing clinical data...")) |
|
|
yield messages, None |
|
|
|
|
|
full_report = [] |
|
|
for i, chunk in enumerate(chunks, 1): |
|
|
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 |
|
|
|
|
|
cleaned = clean_response(response) |
|
|
full_report.append(cleaned) |
|
|
|
|
|
|
|
|
progress_msg = f"✅ Analyzed section {i}/{len(chunks)}" |
|
|
if len(messages) > 2 and "Analyzed section" in messages[-1][1]: |
|
|
messages[-1] = ("assistant", progress_msg) |
|
|
else: |
|
|
messages.append(("assistant", progress_msg)) |
|
|
yield messages, None |
|
|
|
|
|
|
|
|
final_report = "## 🧠 Final Clinical Report\n\n" + "\n\n".join(full_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(final_report) |
|
|
|
|
|
messages.append(("assistant", f"✅ Report generated and saved: {os.path.basename(report_path)}")) |
|
|
messages.append(("assistant", final_report)) |
|
|
yield messages, report_path |
|
|
|
|
|
except Exception as e: |
|
|
messages.append(("assistant", f"❌ Error: {str(e)}")) |
|
|
yield messages, None |
|
|
|
|
|
def create_ui(agent): |
|
|
with gr.Blocks(css=""" |
|
|
body { |
|
|
background: #10141f; |
|
|
color: #ffffff; |
|
|
font-family: 'Inter', sans-serif; |
|
|
margin: 0; |
|
|
padding: 0; |
|
|
} |
|
|
.gradio-container { |
|
|
padding: 30px; |
|
|
width: 100vw; |
|
|
max-width: 100%; |
|
|
border-radius: 0; |
|
|
background-color: #1a1f2e; |
|
|
} |
|
|
.chatbot { |
|
|
background-color: #131720; |
|
|
border-radius: 12px; |
|
|
padding: 20px; |
|
|
height: 600px; |
|
|
overflow-y: auto; |
|
|
border: 1px solid #2c3344; |
|
|
} |
|
|
.gr-button { |
|
|
background: linear-gradient(135deg, #4b4ced, #37b6e9); |
|
|
color: white; |
|
|
font-weight: 500; |
|
|
border: none; |
|
|
padding: 10px 20px; |
|
|
border-radius: 8px; |
|
|
transition: background 0.3s ease; |
|
|
} |
|
|
.gr-button:hover { |
|
|
background: linear-gradient(135deg, #37b6e9, #4b4ced); |
|
|
} |
|
|
.report-content { |
|
|
background-color: #1a1f2e; |
|
|
padding: 15px; |
|
|
border-radius: 8px; |
|
|
margin-top: 10px; |
|
|
border: 1px solid #2c3344; |
|
|
} |
|
|
.bullet-points { |
|
|
margin-left: 20px; |
|
|
} |
|
|
""") as demo: |
|
|
gr.Markdown("""# Clinical Reasoning Assistant |
|
|
Upload clinical Excel records below and click **Analyze** to generate a medical summary. |
|
|
""") |
|
|
|
|
|
chatbot = gr.Chatbot(label="Chatbot", elem_classes="chatbot") |
|
|
file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"]) |
|
|
analyze_btn = gr.Button("Analyze") |
|
|
report_output = gr.File(label="Download Report", visible=False) |
|
|
chatbot_state = gr.State([]) |
|
|
|
|
|
analyze_btn.click( |
|
|
fn=process_final_report, |
|
|
inputs=[file_upload, chatbot_state, gr.State(agent)], |
|
|
outputs=[chatbot, report_output], |
|
|
show_progress="hidden" |
|
|
) |
|
|
|
|
|
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) |
