import gradio as gr import torch import numpy as np from pathlib import Path import re from Model import OmniPathWithInterTaskAttention from transformers import AutoModelForCausalLM, AutoTokenizer import transformers import os # 强制设置 Gradio 为英文环境 os.environ["GRADIO_LOCALE"] = "en" # 设备设置 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") # 预加载模型（避免重复加载） @torch.no_grad() def load_models(): """Preload necessary models""" # 1. Load classification model ckpt_path = "best_model.pth" if not Path(ckpt_path).exists(): raise FileNotFoundError(f"Model file not found: {ckpt_path}") ckpt = torch.load(ckpt_path, map_location=device) label_mappings = ckpt.get('label_mappings', None) if not label_mappings: raise ValueError("The checkpoint is missing label_mappings") ck_cfg = ckpt.get('config', {}) feature_dim = 768 # Adjust according to your actual feature dimension hidden_dim = int(ck_cfg.get('hidden_dim', 256)) dropout = float(ck_cfg.get('dropout', 0.3)) use_inter_task_attention = bool(ck_cfg.get('use_inter_task_attention', True)) inter_task_heads = int(ck_cfg.get('inter_task_heads', 4)) classification_model = OmniPathWithInterTaskAttention( label_mappings=label_mappings, feature_dim=feature_dim, hidden_dim=hidden_dim, dropout=dropout, use_inter_task_attention=use_inter_task_attention, inter_task_heads=inter_task_heads ).to(device) classification_model.load_state_dict(ckpt['model_state_dict'], strict=False) classification_model.eval() # 2. Load text generation model llm_model_name = "Qwen/Qwen3-0.6B" # llm_model_name = "Qwen/QwQ-32B" tokenizer = AutoTokenizer.from_pretrained(llm_model_name) llm_model = AutoModelForCausalLM.from_pretrained( llm_model_name, dtype="auto", device_map="auto" ) return classification_model, llm_model, tokenizer, label_mappings # 预加载模型 classification_model, llm_model, tokenizer, label_mappings = load_models() def analyze_npy_file(npy_file): """Analyze NPY file and return prediction results""" if npy_file is None: return None, "Please upload an NPY file first" try: # Read NPY file arr = np.load(npy_file.name, allow_pickle=False) if not isinstance(arr, np.ndarray) or arr.ndim != 2: return None, "Error: NPY file must be a two-dimensional feature matrix" features = torch.from_numpy(arr).float() # Extract short ID p = Path(npy_file.name) m = re.search(r'(TCGA-[A-Z0-9]{2}-[A-Z0-9]{4})', p.name.upper()) short_id = m.group(1) if m else p.stem[:12] # Inference feat_batch = features.unsqueeze(0).to(device) outputs = classification_model(feat_batch) # Decode results pred_names, pred_scores = {}, {} for task_name, logits in outputs.items(): probs = torch.softmax(logits[0], dim=-1) idx = int(torch.argmax(probs).item()) classes = label_mappings[task_name]['classes'] class_name = classes[idx] if 0 <= idx < len(classes) else str(idx) pred_names[task_name] = class_name pred_scores[task_name] = float(probs[idx].item()) # Format results results_text = f"Patient ID: {short_id}\n\nPrediction Results:\n" for task, name in pred_names.items(): results_text += f"- {task}: {name} (Confidence: {pred_scores.get(task, 0.0):.3f})\n" return {"pred_names": pred_names, "pred_scores": pred_scores, "patient_id": short_id}, results_text except Exception as e: return None, f"An error occurred during processing: {str(e)}" def generate_response(message, chat_history, analysis_results): """Generate response based on user message and analysis results""" if analysis_results is None: return "Please upload an NPY file first to analyze the patient data.", chat_history pred_names = analysis_results["pred_names"] pred_scores = analysis_results["pred_scores"] patient_id = analysis_results["patient_id"] # Build context from analysis results context = f"Patient {patient_id} analysis results:\n" for task, name in pred_names.items(): context += f"- {task}: {name} (confidence: {pred_scores.get(task, 0.0):.3f})\n" # Build prompt based on user message if "diagnosis" in message.lower() or "result" in message.lower(): prompt = f"{context}\nBased on the above analysis results, provide a detailed diagnosis summary and interpretation." elif "treatment" in message.lower() or "therapy" in message.lower(): prompt = f"{context}\nBased on the diagnosis, suggest appropriate treatment options and considerations." elif "prognosis" in message.lower() or "outlook" in message.lower(): prompt = f"{context}\nDiscuss the prognosis and potential outcomes for this patient." elif "stage" in message.lower(): prompt = f"{context}\nExplain the staging information and its clinical implications." elif "histology" in message.lower() or "type" in message.lower(): prompt = f"{context}\nDescribe the histological characteristics and their significance." else: prompt = f"{context}\nUser question: {message}\nPlease provide a helpful response based on the analysis results." try: # Generate response using LLM messages = [{"role": "user", "content": prompt}] text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False ) model_inputs = tokenizer([text], return_tensors="pt").to(llm_model.device) generated_ids = llm_model.generate( **model_inputs, max_new_tokens=32768, do_sample=True, temperature=0.7, ) output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() try: index = len(output_ids) - output_ids[::-1].index(151668) except ValueError: index = 0 response = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n") # Add to chat history chat_history.append((message, response)) return "", chat_history except Exception as e: error_msg = f"Error generating response: {str(e)}" chat_history.append((message, error_msg)) return "", chat_history def upload_file(npy_file, chat_history, analysis_results): """Handle file upload and initial analysis""" if npy_file is None: return chat_history, analysis_results, "Please select a file to upload" new_analysis_results, results_text = analyze_npy_file(npy_file) if new_analysis_results is None: return chat_history, analysis_results, results_text # Add analysis results to chat chat_history.append(("System", f"File uploaded and analyzed successfully!\n{results_text}")) chat_history.append(("System", "You can now ask questions about this patient's diagnosis, treatment options, prognosis, etc.")) return chat_history, new_analysis_results, "Analysis completed successfully!" def example_click(example): """Handle example question click""" return example # Create conversational interface with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown(""" # 🏥 Medical Pathology Diagnostic Chat Assistant Upload a pathology NPY file and chat with the AI assistant about the diagnosis, treatment options, prognosis, and more. """) # Store analysis results in session state analysis_results = gr.State(value=None) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Upload Patient Data") file_input = gr.File( label="Upload NPY Feature File", file_types=[".npy"], type="filepath" ) upload_btn = gr.Button("Upload & Analyze", variant="primary") status_output = gr.Textbox( label="Status", lines=2, interactive=False ) with gr.Column(scale=2): gr.Markdown("### Chat with Medical Assistant") chatbot = gr.Chatbot( label="Conversation", height=400 ) with gr.Row(): msg = gr.Textbox( label="Your Question", placeholder="Ask about diagnosis, treatment, prognosis...", lines=2, scale=4 ) send_btn = gr.Button("Send", variant="primary", scale=1) with gr.Row(): clear_btn = gr.Button("Clear Chat") gr.Markdown("### Suggested Questions") examples = gr.Examples( examples=[ "What is the diagnosis?", "What treatment options are available?", "What is the prognosis?", "Explain the staging information", "Describe the histological findings" ], inputs=msg, # 将示例应用到消息输入框 fn=example_click, # 点击示例时的处理函数 outputs=msg, # 输出到消息输入框 label="Click a question to use it" ) # Event handlers upload_btn.click( upload_file, inputs=[file_input, chatbot, analysis_results], outputs=[chatbot, analysis_results, status_output] ) send_btn.click( generate_response, inputs=[msg, chatbot, analysis_results], outputs=[msg, chatbot] ) msg.submit( generate_response, inputs=[msg, chatbot, analysis_results], outputs=[msg, chatbot] ) clear_btn.click( lambda: ([], None, "Chat cleared"), inputs=[], outputs=[chatbot, analysis_results, status_output] ) if __name__ == "__main__": demo.launch(share=True)