jinge13288/RiskAgent
0
1import os2from typing import Optional, List, Dict3from contextlib import asynccontextmanager4import re5import json6 7from fastapi import FastAPI, HTTPException, status8from fastapi.middleware.cors import CORSMiddleware9from pydantic import BaseModel10from transformers import AutoTokenizer, AutoModelForCausalLM11import torch12import gradio as gr13 14class MedicalReport(BaseModel):15 text: str16 17class ReportResponse(BaseModel):18 assessment: str19 20class MedicalAssessmentModel:21 def __init__(self):22 # Initialize model and tokenizer23 model_name = "meta-llama/Llama-2-7b-chat-hf" # or any other model you prefer24 self.tokenizer = AutoTokenizer.from_pretrained(model_name)25 self.model = AutoModelForCausalLM.from_pretrained(26 model_name,27 torch_dtype=torch.float16,28 device_map="auto"29 )30 31 def generate_response(self, messages: List[Dict]) -> str:32 # Combine messages into a single prompt33 prompt = ""34 for msg in messages:35 role = msg['role']36 content = msg['content']37 prompt += f"{role}: {content}\n"38 39 inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)40 41 with torch.no_grad():42 outputs = self.model.generate(43 **inputs,44 max_new_tokens=512,45 temperature=0.7,46 do_sample=True,47 top_p=0.9,48 num_return_sequences=1,49 )50 51 response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)52 # Extract only the generated part53 response = response[len(self.tokenizer.decode(inputs['input_ids'][0], skip_special_tokens=True)):]54 return response.strip()55 56 def run_env1(self, patient_text: str) -> str:57 """Tool Selection Stage"""58 messages = [59 {60 "role": "system",61 "content": "You are a medical professional expert in selecting appropriate clinical risk assessment tools."62 },63 {64 "role": "user",65 "content": f"""Based on the patient's discharge summary, identify potential disease risks and assessment needs.66 67Patient Information:68{patient_text}69 70Please analyze:711. Primary health concerns722. Risk factors identified733. Potential complications744. Areas requiring risk assessment"""75 }76 ]77 return self.generate_response(messages)78 79 def run_env2(self, patient_text: str, env1_output: str) -> str:80 """Parameter Extraction Stage"""81 messages = [82 {83 "role": "system",84 "content": "You are a medical professional expert in extracting clinical parameters from patient records."85 },86 {87 "role": "user",88 "content": f"""Extract relevant clinical parameters from the patient's information.89 90Patient Information:91{patient_text}92 93Previous Analysis:94{env1_output}95 96Please provide:971. Key vital signs982. Relevant lab values993. Clinical findings1004. Risk factors identified"""101 }102 ]103 return self.generate_response(messages)104 105 def run_env3(self, patient_text: str, env1_output: str, env2_output: str) -> str:106 """Risk Interpretation Stage"""107 messages = [108 {109 "role": "system",110 "content": "You are a medical expert specialized in clinical risk assessment and interpretation."111 },112 {113 "role": "user",114 "content": f"""Interpret the identified risks and clinical parameters.115 116Patient Information:117{patient_text}118 119Risk Analysis:120{env1_output}121 122Clinical Parameters:123{env2_output}124 125Please provide:1261. Risk level assessment for each identified condition1272. Clinical significance of findings1283. Interaction between different risk factors1294. Severity assessment"""130 }131 ]132 return self.generate_response(messages)133 134 def run_env4(self, patient_text: str, env1_output: str, env2_output: str, env3_output: str) -> str:135 """Final Assessment Stage"""136 messages = [137 {138 "role": "system",139 "content": "You are a medical expert specialized in comprehensive risk assessment and patient care planning."140 },141 {142 "role": "user",143 "content": f"""Based on all previous analyses, provide a comprehensive assessment of the patient's disease risks.144 145Patient Information:146{patient_text}147 148Previous Analyses:149Risk Identification: {env1_output}150Parameter Analysis: {env2_output}151Risk Interpretation: {env3_output}152 153Please provide:1541. Summary of significant disease risks identified1552. Overall risk assessment1563. Key areas of concern1574. Recommended monitoring or preventive measures1585. Suggestions for risk mitigation159 160Format the response in clear sections with headers."""161 }162 ]163 return self.generate_response(messages)164 165 def process_report(self, patient_text: str) -> str:166 """Process the entire pipeline and return ENV4 output"""167 try:168 # Run all environments sequentially169 env1_output = self.run_env1(patient_text)170 env2_output = self.run_env2(patient_text, env1_output)171 env3_output = self.run_env3(patient_text, env1_output, env2_output)172 env4_output = self.run_env4(patient_text, env1_output, env2_output, env3_output)173 174 return env4_output175 except Exception as e:176 return f"Error in processing: {str(e)}"177 178def create_gradio_interface():179 model = MedicalAssessmentModel()180 181 def analyze_text(text):182 return model.process_report(text)183 184 iface = gr.Interface(185 fn=analyze_text,186 inputs=gr.Textbox(187 lines=10, 188 placeholder="Enter patient medical report here...",189 label="Medical Report"190 ),191 outputs=gr.Textbox(192 lines=15,193 label="Risk Assessment Report"194 ),195 title="Medical Report Risk Assessment",196 description="Enter a medical report to get a comprehensive risk assessment. The system will analyze the report through multiple stages and provide a final assessment.",197 examples=[198 ["Patient was admitted with chest pain and shortness of breath. History of hypertension and diabetes. BP 160/95, HR 98. Recent smoker with 30 pack-year history."],199 ["83-year-old female presents with confusion and fever. Recent fall at home. History of osteoporosis and mild cognitive impairment. Lives alone. Temperature 38.5C, BP 135/85."]200 ]201 )202 return iface203 204if __name__ == "__main__":205 iface = create_gradio_interface()206 iface.launch(server_name="0.0.0.0", server_port=7860)