app.py

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)