File size: 9,802 Bytes
1bb8be7
dae38a2
 
 
 
1bb8be7
e24be23
dae38a2
 
 
e24be23
 
1da2cfd
 
dae38a2
1da2cfd
e24be23
 
 
 
1da2cfd
e24be23
1da2cfd
e24be23
 
 
1ebbef1
dae38a2
 
e24be23
 
1ebbef1
e24be23
1da2cfd
 
 
 
 
 
dae38a2
 
 
1da2cfd
 
 
 
 
dae38a2
 
 
e24be23
dae38a2
 
 
1da2cfd
 
 
 
 
 
 
 
67f4d88
1da2cfd
 
 
 
e24be23
dae38a2
 
 
e24be23
dae38a2
 
 
1da2cfd
 
1ebbef1
1da2cfd
722c891
1da2cfd
722c891
1da2cfd
 
722c891
dae38a2
 
 
 
 
1da2cfd
 
722c891
dae38a2
 
 
e24be23
 
dae38a2
722c891
dae38a2
1da2cfd
 
 
 
 
 
 
 
722c891
1ebbef1
1da2cfd
 
1ebbef1
 
1da2cfd
 
e24be23
 
 
 
 
 
 
 
1da2cfd
 
e24be23
 
 
 
 
722c891
e24be23
 
 
 
dae38a2
 
1da2cfd
 
dae38a2
722c891
1ebbef1
 
 
 
 
1da2cfd
 
dae38a2
1ebbef1
dae38a2
1da2cfd
e24be23
dae38a2
67f4d88
 
1ebbef1
722c891
1da2cfd
1ebbef1
1da2cfd
 
722c891
1ebbef1
 
 
1da2cfd
67f4d88
1da2cfd
 
 
 
 
1ebbef1
1da2cfd
1ebbef1
1da2cfd
 
 
 
dae38a2
722c891
1da2cfd
e24be23
dae38a2
1da2cfd
 
 
 
 
 
1ebbef1
722c891
67f4d88
1ebbef1
1da2cfd
 
 
 
 
 
1ebbef1
1da2cfd
1ebbef1
 
dae38a2
1da2cfd
67f4d88
1ebbef1
dae38a2
1da2cfd
1ebbef1
 
 
dae38a2
 
1da2cfd
 
 
 
1bb8be7
dae38a2
e24be23
 
1da2cfd
e24be23
 
 
722c891
e24be23
 
 
 
1ebbef1
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
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

# 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 = os.path.join(base_dir, "reports")

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.environ.update({
    "TRANSFORMERS_CACHE": model_cache_dir,
    "HF_HOME": model_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 (after tools finish)")

        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:
                    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')]
                        results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
                        extracted_data = "\n".join(results)
                        file_hash_value = file_hash(files[0].name)

                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.append(chunk)
                    elif isinstance(chunk, list):
                        response.extend([c.content for c in chunk if hasattr(c, 'content')])

                    if len(response) % 3 == 0:
                        history[-1] = {"role": "assistant", "content": "".join(response).strip()}
                        yield history, None

                final_output = "".join(response).strip()
                if not final_output:
                    final_output = "No clear oversights identified. Recommend comprehensive review."
                if not final_output.startswith(("1.", "-", "*", "#")):
                    final_output = "• " + final_output.replace("\n", "\n• ")
                history[-1] = {"role": "assistant", "content": final_output}

                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
                return history, report_path if os.path.exists(report_path) else None

            except Exception as e:
                history.append({"role": "assistant", "content": f"❌ Analysis failed: {str(e)}"})
                return 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,
        share=True
    )