Direction9/model/extract_representations.py

#!/usr/bin/env python
"""
从预训练模型提取文本表示

使用现有的模型权重，为每个领域的文本提取最后一层隐藏状态，
计算 mean pooling 后保存为 H ∈ R^{d×N} 矩阵。
"""

import os
# 设置环境变量以减少显存碎片化
os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"

import gc
import torch
import json
from pathlib import Path
from typing import List, Dict
from transformers import AutoTokenizer, AutoModel
import warnings

warnings.filterwarnings("ignore")


def load_texts(filepath: str) -> List[str]:
    """从 JSONL 文件加载文本"""
    texts = []
    with open(filepath, "r", encoding="utf-8") as f:
        for line in f:
            record = json.loads(line.strip())
            texts.append(record.get("text", ""))
    return texts


def extract_representations(
    model_path: str,
    texts: List[str],
    batch_size: int = 4,  # 减小 batch_size 以避免 OOM
    max_length: int = 512
) -> torch.Tensor:
    """
    从文本列表提取模型表示

    Args:
        model_path: 模型路径（本地目录）
        texts: 文本列表
        batch_size: batch 大小
        max_length: 最大序列长度

    Returns:
        H: 表示矩阵，shape = (d, N)
    """
    print(f"  加载模型：{model_path}")

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

    # 修复：某些模型（如 Mistral）缺少 pad_token，需要手动设置
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        print(f"    设置 pad_token = eos_token: {repr(tokenizer.pad_token)}")

    model = AutoModel.from_pretrained(
        model_path,
        torch_dtype=torch.float32,
        device_map="auto",
        trust_remote_code=True,
        output_hidden_states=True
    )
    model.eval()

    print(f"  模型设备：{model.device}")
    print(f"  文本数量：{len(texts)}")

    all_embeddings = []

    with torch.no_grad():
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]

            # Tokenize
            inputs = tokenizer(
                batch_texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=max_length
            )
            inputs = inputs.to(model.device)

            # Forward pass
            outputs = model(**inputs)

            # Mean pooling over sequence length
            # hidden_states: (batch, seq_len, hidden_dim)
            hidden = outputs.last_hidden_state
            attention_mask = inputs["attention_mask"].unsqueeze(-1).float()

            # 加权平均（考虑 padding）
            sum_embeddings = (hidden * attention_mask).sum(dim=1)
            valid_tokens = attention_mask.sum(dim=1)
            mean_embeddings = sum_embeddings / (valid_tokens + 1e-9)

            all_embeddings.append(mean_embeddings.cpu())

            print(f"    进度：{min(i + batch_size, len(texts))}/{len(texts)}")

    # 拼接所有 embedding: (N, d)
    embeddings = torch.cat(all_embeddings, dim=0)

    # 转置为 (d, N) 格式
    H = embeddings.T

    # 清理 GPU 内存
    del model, tokenizer
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return H


def save_representation(
    H: torch.Tensor,
    model_name: str,
    domain: str,
    output_dir: str = "model/representations"
) -> str:
    """保存表示矩阵"""
    output_path = Path(output_dir) / f"{model_name}_{domain}.pt"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    torch.save(H, output_path)

    # 保存元数据
    meta_path = output_path.with_suffix(".json")
    meta = {
        "model_name": model_name,
        "domain": domain,
        "d": int(H.shape[0]),
        "N": int(H.shape[1]),
        "dtype": str(H.dtype)
    }
    with open(meta_path, "w") as f:
        json.dump(meta, f, indent=2)

    return str(output_path)


def get_model_name(model_path: str) -> str:
    """从路径提取模型简称"""
    path = Path(model_path)
    name = path.name.lower()

    # 简化命名
    if "qwen" in name:
        return "qwen_7b"
    elif "mistral" in name:
        return "mistral_7b"
    elif "llama3" in name or "llama-3" in name:
        return "llama3_8b"
    elif "gemma" in name:
        return "gemma_9b"
    elif "bert" in name:
        return "bert"
    elif "gpt2" in name:
        return "gpt2"
    else:
        return path.name.replace("-", "_")


def main():
    print("=" * 60)
    print("模型表示提取")
    print("=" * 60)
    print()

    # 可用模型
    models = [
        "model/weights/Qwen2.5-7B-Instruct",
        "model/weights/Mistral-7B-v0.3",
        "model/weights/llama3-8b",
    ]

    # 可用领域（6 个原始领域 + 5 个 GLUE 领域）
    domains = [
        # 原始领域
        ("news_en", "database/corpus/news_en/texts.jsonl"),
        ("news_zh", "database/corpus/news_zh/texts.jsonl"),
        ("academic", "database/corpus/academic/texts.jsonl"),
        ("code", "database/corpus/code/texts.jsonl"),
        ("dialogue", "database/corpus/dialogue/texts.jsonl"),
        ("literature", "database/corpus/literature/texts.jsonl"),
        # GLUE 领域
        ("glue_mnli", "database/corpus/glue_mnli/texts.jsonl"),
        ("glue_qnli", "database/corpus/glue_qnli/texts.jsonl"),
        ("glue_sst2", "database/corpus/glue_sst2/texts.jsonl"),
        ("glue_stsb", "database/corpus/glue_stsb/texts.jsonl"),
        ("glue_cola", "database/corpus/glue_cola/texts.jsonl"),
    ]

    # 检查哪些文本文件存在
    available_domains = []
    for domain, path in domains:
        if Path(path).exists():
            available_domains.append((domain, path))
        else:
            print(f"跳过（文件不存在）：{path}")

    if not available_domains:
        print("错误：没有找到任何文本语料文件！")
        return

    print(f"可用领域：{[d[0] for d in available_domains]}")
    print(f"模型数量：{len(models)}")
    print()

    # 对每个模型 × 领域组合提取表示
    for model_path in models:
        if not Path(model_path).exists():
            print(f"跳过（模型不存在）：{model_path}")
            continue

        model_name = get_model_name(model_path)
        print(f"\n[模型：{model_name}]")

        for domain, texts_path in available_domains:
            print(f"\n  [{domain}]")

            # 加载文本
            texts = load_texts(texts_path)
            print(f"    加载了 {len(texts)} 条文本")

            if len(texts) == 0:
                print("    跳过：没有文本")
                continue

            # 提取表示
            H = extract_representations(model_path, texts)

            print(f"    表示矩阵形状：{H.shape}")
            print(f"    dtype: {H.dtype}")

            # 保存
            output_path = save_representation(H, model_name, domain)
            print(f"    已保存：{output_path}")

            # 每个领域结束后清理 GPU 缓存
            del H
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    print()
    print("=" * 60)
    print("提取完成!")
    print("=" * 60)
    print()
    print("下一步：运行 MVP 实验")
    print("  python -m experiments.verify --mode mvp")


if __name__ == "__main__":
    main()