LLM_new.py

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()