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 tempfile import os # 设备设置 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"使用设备: {device}") # 预加载模型（避免重复加载） @torch.no_grad() def load_models(): """预加载必要的模型""" # 1. 加载分类模型 ckpt_path = "best_model.pth" if not Path(ckpt_path).exists(): raise FileNotFoundError(f"找不到模型文件: {ckpt_path}") ckpt = torch.load(ckpt_path, map_location=device) label_mappings = ckpt.get('label_mappings', None) if not label_mappings: raise ValueError("checkpoint 中缺少 label_mappings") ck_cfg = ckpt.get('config', {}) feature_dim = 512 # 根据你的实际特征维度调整 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. 加载文本生成模型 llm_model_name = "Qwen/Qwen3-0.6B" tokenizer = AutoTokenizer.from_pretrained(llm_model_name) llm_model = AutoModelForCausalLM.from_pretrained( llm_model_name, torch_dtype="auto", device_map="auto" ) return classification_model, llm_model, tokenizer, label_mappings # 预加载模型 classification_model, llm_model, tokenizer, label_mappings = load_models() def build_prompt(pred_names, pred_scores): """构建提示词""" def get_pred(task_name): name = pred_names.get(task_name, "N/A") score = pred_scores.get(task_name, 0.0) return f"{name} (confidence: {score:.1%})" cancer_type = get_pred('cancer_type') pathologic_stage = get_pred('pathologic_stage') clinical_stage = get_pred('clinical_stage') histological_type = get_pred('histological_type') prompt = ( "You are a professional medical report generator. " "Based on the patient's pathological classification and diagnostic model results, " f"the cancer_type is {cancer_type}, " f"pathologic_stage is {pathologic_stage}, " f"clinical_stage is {clinical_stage}, " f"and histological_type is {histological_type}. " "Please write a concise English summary describing the diagnosis, staging interpretation, and general clinical implications as a short paragraph. " "Avoid placeholders and avoid repeating words." ) return prompt def generate_description(predictions, confidences): """生成描述""" prompt = build_prompt(predictions, confidences) 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=500, # 减少token数量以加快生成速度 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 content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n") return {"prompt": prompt, "description": content} def process_npy_file(npy_file): """处理上传的NPY文件""" if npy_file is None: return "请先上传NPY文件", "", "" try: # 读取NPY文件 arr = np.load(npy_file.name, allow_pickle=False) if not isinstance(arr, np.ndarray) or arr.ndim != 2: return "错误: NPY文件必须是二维特征矩阵", "", "" features = torch.from_numpy(arr).float() # 提取短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] # 推理 feat_batch = features.unsqueeze(0).to(device) outputs = classification_model(feat_batch) # 解码结果 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()) # 生成描述 desc_result = generate_description(pred_names, pred_scores) # 格式化结果显示 results_text = f"患者ID: {short_id}\n\n预测结果:\n" for task, name in pred_names.items(): results_text += f"- {task}: {name} (置信度: {pred_scores.get(task, 0.0):.3f})\n" return results_text, desc_result['prompt'], desc_result['description'] except Exception as e: return f"处理过程中出现错误: {str(e)}", "", "" # 创建Gradio界面 with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown(""" # 🏥 医学病理诊断分析系统 上传病理特征NPY文件，获取AI辅助诊断结果和医学描述。 """) with gr.Row(): with gr.Column(): file_input = gr.File( label="上传NPY特征文件", file_types=[".npy"], type="filepath" ) analyze_btn = gr.Button("开始分析", variant="primary") with gr.Column(): results_output = gr.Textbox( label="分析结果", lines=10, max_lines=20 ) with gr.Row(): prompt_output = gr.Textbox( label="提示词 (Prompt)", lines=4, max_lines=6 ) with gr.Row(): description_output = gr.Textbox( label="AI生成描述", lines=6, max_lines=10 ) # 示例文件 gr.Examples( examples=[["example.npy"]], # 你需要提供一个示例NPY文件 inputs=file_input, label="点击使用示例文件" ) analyze_btn.click( fn=process_npy_file, inputs=file_input, outputs=[results_output, prompt_output, description_output] ) if __name__ == "__main__": demo.launch(share=True)