Upload LLM_new.py

LLM_new.py

@@ -0,0 +1,243 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from typing import Dict
+from typing import Dict
+import torch
+from pathlib import Path
+import numpy as np
+import re
+from Model import OmniPathWithInterTaskAttention
+
+def build_prompt(pred_names: Dict[str, str], pred_scores: Dict[str, float]) -> str:
+    """
+    根据分类结果，构建一个稳定、具医学上下文的提示，用于LLM生成简洁英文描述。
+    """
+    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
+# prepare the model input
+
+def build_description(predictions: Dict[str, str], confidences: Dict[str, float],
+                      model_name: str = "Qwen/Qwen3-0.6B", max_new_tokens: int = 32768) -> Dict[str, str]:
+    """
+    Compose the prompt and obtain a concise description using a model when available,
+    otherwise fall back to a deterministic template. Output length is controlled by max_new_tokens.
+    """
+    prompt = build_prompt(predictions, confidences)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto"
+    )
+    messages = [
+        {"role": "user", "content": prompt}
+    ]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=False # Switches between thinking and non-thinking modes. Default is True.
+    )
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=32768
+    )
+    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
+    try:
+    # rindex finding 151668 (</think>)
+        index = len(output_ids) - output_ids[::-1].index(151668)
+    except ValueError:
+        index = 0
+    thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
+    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+
+    print("thinking content:", thinking_content)
+    print("content:", content)
+
+    return {"prompt": prompt, "description": content}
+
+@torch.no_grad()
+def llm_single_patient_test(config, device):
+    """
+    使用已训练好的 checkpoint，对单个小NPY患者进行推理测试（不依赖大NPY的数据集/数据加载器）。
+
+    要求 config 提供：
+    - npy_path: 小NPY文件路径（二维矩阵，形状 [num_tiles, feature_dim]）
+    - checkpoint_path: 训练得到的 best_model.pth 路径（包含 label_mappings 与模型权重）
+
+    返回: {
+      'short_id': str,
+      'pred_names': Dict[task, class_name],
+      'pred_scores': Dict[task, float],
+      'raw_logits': Dict[task, torch.Tensor]
+    }
+    """
+    # 1) 读取小NPY（直接在此函数内完成）
+    npy_path = config['npy_path'] if isinstance(config, dict) else getattr(config, 'npy_path', None)
+    if not npy_path:
+        raise ValueError("config 中未提供 'npy_path'")
+    p = Path(npy_path)
+    if not p.exists():
+        raise FileNotFoundError(f"找不到npy文件: {npy_path}")
+    # 提取短ID（文件名前12位 TCGA-XX-XXXX）
+    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]
+    arr = np.load(str(p), allow_pickle=False)
+    if not isinstance(arr, np.ndarray) or arr.ndim != 2:
+        raise ValueError(
+            f"npy 内容必须是二维特征矩阵 (tiles, dim)，实际: type={type(arr)}, shape={getattr(arr, 'shape', None)}"
+        )
+    features = torch.from_numpy(arr).float()  # [N, D]
+
+    # 2) 读取checkpoint
+    ckpt_path = config['checkpoint_path'] if isinstance(config, dict) else getattr(config, 'checkpoint_path', None)
+    if not ckpt_path:
+        raise ValueError("config 中未提供 'checkpoint_path'")
+    ckpt = torch.load(ckpt_path, map_location=device)
+
+    # 3) 构建模型（用checkpoint内保存的label_mappings与config参数）
+    label_mappings = ckpt.get('label_mappings', None)
+    if not label_mappings:
+        raise ValueError("checkpoint 中缺少 label_mappings，无法构建模型")
+
+    ck_cfg = ckpt.get('config', {}) if isinstance(ckpt.get('config', {}), dict) else {}
+    feature_dim = int(features.shape[1])
+    hidden_dim = int(ck_cfg.get('hidden_dim', 256))
+    dropout = float(ck_cfg.get('dropout', 0.3)) if 'dropout' in ck_cfg else 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))
+
+    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)
+    model.load_state_dict(ckpt['model_state_dict'], strict=False)
+    model.eval()
+
+    # 4) 前向推理
+    feat_batch = features.unsqueeze(0).to(device)  # [1, N, D]
+    outputs = model(feat_batch)  # {task: [1, num_classes]}
+
+    # 5) 解码到类别名称与置信度
+    pred_names, pred_scores, raw_logits = {}, {}, {}  # 为每个任务分别存放类别名、置信度和原始logits
+    for task_name, logits in outputs.items():  # 遍历各任务的输出（形如 [1, num_classes] 的logits）
+        probs = torch.softmax(logits[0], dim=-1)  # 对单样本的logits做softmax，得到每个类别的概率分布
+        idx = int(torch.argmax(probs).item())  # 取概率最大的类别索引，作为预测类别
+        # 映射 idx -> class name
+        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())  # 记录该任务的预测置信度（最大概率）
+        raw_logits[task_name] = logits[0].detach().cpu()  # 保存原始logits（去梯度并搬到CPU，便于后续分析/可视化）
+
+    return {
+        'short_id': short_id,
+        'pred_names': pred_names,
+        'pred_scores': pred_scores,
+        'raw_logits': raw_logits
+    }
+
+
+def main():
+    # 设备
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"🖥️  使用设备: {device}")
+
+    # 配置（请修改为你的实际路径）
+    config = {
+        # 必填：小NPY路径（二维矩阵）
+        'npy_path': 'TCGA-56-8304-01Z-00-DX1.F7A7975D-C8AB-49C0-B9CD-18CFD01A0655.svs.npy',
+        # 必填：训练阶段保存的最佳checkpoint
+        'checkpoint_path': 'best_model.pth',
+        # 可选：文本描述生成器配置
+        'model_name': "Qwen/Qwen3-0.6B",
+        'max_new_tokens': 32768,
+    }
+
+    # 基础检查
+    if not Path(config['npy_path']).exists():
+        raise FileNotFoundError(f"npy_path 不存在: {config['npy_path']}")
+    if not Path(config['checkpoint_path']).exists():
+        raise FileNotFoundError(f"checkpoint_path 不存在: {config['checkpoint_path']}")
+
+    # 1) 单病人推理
+    result = llm_single_patient_test(config, device)
+    short_id = result['short_id']
+    pred_names = result['pred_names']
+    pred_scores = result['pred_scores']
+
+    print("\n预测结果（按任务）:")
+    for task, name in pred_names.items():
+        print(f"- {task}: {name} (conf {pred_scores.get(task, 0.0):.3f})")
+
+    # 2) 生成文字描述
+    desc = build_description(
+        predictions=pred_names,
+        confidences=pred_scores,
+        model_name=config.get('model_name', "Qwen/Qwen3-0.6B"),
+        max_new_tokens=int(config.get('max_new_tokens', 32768))
+    )
+
+    print("\n=== 自动生成的英文描述 ===")
+    print(desc['description'])
+
+    # 2.1 保存完整结果到文本文件
+    full_txt_path = f"{short_id}_description.txt"
+    try:
+        with open(full_txt_path, 'w', encoding='utf-8') as f:
+            f.write("=== PROMPT ===\n")
+            f.write(desc['prompt'])
+            f.write("\n\n=== DESCRIPTION (full paragraph) ===\n")
+            f.write(desc['description'])
+        print(f"\n📝 Full text saved to: {full_txt_path}")
+    except Exception as e:
+        print(f"Could not save description to file: {e}")
+
+    # === 3) 可视化显示 ===
+    import textwrap
+
+    # 3.1 控制台彩色输出
+    print("\n" + "="*80)
+    print("\033[1;34m🧠 PROMPT:\033[0m\n")
+    wrapped_prompt = textwrap.fill(desc['prompt'], width=100)
+    print(f"\033[0;37m{wrapped_prompt}\033[0m")
+
+    print("\n\033[1;32m💬 DESCRIPTION:\033[0m\n")
+    wrapped_desc = textwrap.fill(desc['description'], width=100)
+    print(f"\033[0;37m{wrapped_desc}\033[0m")
+    print("="*80 + "\n")
+    
+    return {
+        'short_id': short_id,
+        'predictions': pred_names,
+        'confidences': pred_scores,
+        'description': desc['description'],
+        'prompt': desc['prompt'],
+    }
+
+if __name__ == "__main__":
+    main()