|
|
import sys |
|
|
import os |
|
|
import pandas as pd |
|
|
import json |
|
|
import gradio as gr |
|
|
from typing import List, Tuple, Union, Generator |
|
|
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_obj: Union[str, os.PathLike, 'file']) -> str: |
|
|
all_text = [] |
|
|
try: |
|
|
xls = pd.ExcelFile(file_obj) |
|
|
except Exception as e: |
|
|
raise ValueError(f"❌ Error reading Excel file: {e}") |
|
|
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, max_tokens: int = MAX_CHUNK_TOKENS, max_chunks: int = 30) -> 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)) |
|
|
if len(chunks) >= max_chunks: |
|
|
break |
|
|
curr_chunk, curr_tokens = [line], t |
|
|
else: |
|
|
curr_chunk.append(line) |
|
|
curr_tokens += t |
|
|
if curr_chunk and len(chunks) < max_chunks: |
|
|
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 stream_report(agent, file: Union[str, 'file'], full_output: str) -> Generator: |
|
|
accumulated_text = "" |
|
|
try: |
|
|
if file is None: |
|
|
yield "❌ Please upload a valid Excel file.", None, "" |
|
|
return |
|
|
|
|
|
if hasattr(file, "read"): |
|
|
text = extract_text_from_excel(file) |
|
|
elif isinstance(file, str) and os.path.exists(file): |
|
|
text = extract_text_from_excel(file) |
|
|
else: |
|
|
raise ValueError("❌ Invalid or missing file.") |
|
|
|
|
|
chunks = split_text_into_chunks(text) |
|
|
|
|
|
for i, chunk in enumerate(chunks): |
|
|
prompt = build_prompt_from_text(chunk) |
|
|
partial = "" |
|
|
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): |
|
|
partial += res |
|
|
elif hasattr(res, "content"): |
|
|
partial += res.content |
|
|
cleaned = clean_response(partial) |
|
|
accumulated_text += f"\n\n📄 **Chunk {i+1}**:\n{cleaned}" |
|
|
yield accumulated_text, None, "" |
|
|
|
|
|
summary_prompt = f"Summarize this analysis in a final structured report:\n\n" + accumulated_text |
|
|
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) |
|
|
accumulated_text += f"\n\n📊 **Final Summary**:\n{cleaned}" |
|
|
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}") |
|
|
|
|
|
yield accumulated_text, report_path, cleaned |
|
|
|
|
|
except Exception as e: |
|
|
yield f"❌ Error: {str(e)}", 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; |
|
|
} |
|
|
.output-markdown { |
|
|
background-color: #131720; |
|
|
border-radius: 12px; |
|
|
padding: 20px; |
|
|
min-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); |
|
|
} |
|
|
""") as demo: |
|
|
gr.Markdown("""# 🧠 Clinical Reasoning Assistant |
|
|
Upload clinical Excel records below and click **Analyze** to generate a medical summary. |
|
|
""") |
|
|
file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"]) |
|
|
analyze_btn = gr.Button("Analyze") |
|
|
report_output_markdown = gr.Markdown(elem_classes="output-markdown") |
|
|
report_file = gr.File(label="Download Report", visible=False) |
|
|
full_output = gr.State(value="") |
|
|
|
|
|
analyze_btn.click( |
|
|
fn=lambda file, state: stream_report(agent, file, state), |
|
|
inputs=[file_upload, full_output], |
|
|
outputs=[report_output_markdown, report_file, full_output] |
|
|
) |
|
|
|
|
|
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) |
|
|
|
