Direction9/experiments/verify.py

"""
LLM 有效秩谱收敛验证实验主流程

验证文档《基于 SVCCA 的数据集和模型评价.md》附录 C 中的实验设计

核心研究问题:
    Q1: 不同 LLM 在相同文本集上的谱有效秩是否趋同？
    Q2: 不仅有效秩，特征值分布的整体形状是否收敛？
    Q3: 收敛程度是否随文本领域（新闻/代码/对话/学术）变化？
    Q4: 模型规模越大，谱越接近某个"吸引子"？

实验类型:
    - MVP 实验：3 个模型 × 2 个领域，快速验证
    - 完整实验：7 个模型 × 6 个领域，全面验证
    - 规模效应实验：同家族不同规模模型对比
"""

import torch
import numpy as np
import json
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, field
from datetime import datetime

# 导入核心模块
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))

from model.spectrum import (
    compute_gram_spectrum,
    compute_spectrum_array,
    wasserstein1_distance,
    compute_spectrum_distance_matrix,
    spectrum_summary,
    effective_rank_from_eigvals,
)
from model.extractor import (
    RepresentationLoader,
    load_H,
    load_all_H,
)
from database.corpus import (
    CorpusLoader,
    load_texts,
    load_all_texts,
    list_domains,
)


@dataclass
class ExperimentConfig:
    """实验配置"""
    # 模型配置
    model_names: List[str] = field(default_factory=list)

    # 数据集配置
    domains: List[str] = field(default_factory=list)

    # 实验参数
    limit_per_domain: Optional[int] = None  # 每个领域限制多少文本

    # 计算参数
    normalize_method: str = "per_sample_l2"

    # 输出配置
    output_dir: str = "experiments/output"

    @classmethod
    def mvp_config(cls) -> "ExperimentConfig":
        """MVP 实验配置（最小可行实验）"""
        return cls(
            model_names=["qwen_7b", "mistral_7b", "llama3_8b"],
            domains=["news_en", "news_zh", "academic", "code", "dialogue", "literature"],
            limit_per_domain=500,
            output_dir="experiments/output/mvp",
        )

    @classmethod
    def glue_config(cls, limit_per_domain: Optional[int] = 5000) -> "ExperimentConfig":
        """GLUE Benchmark 完整数据集配置"""
        return cls(
            model_names=["qwen_7b", "mistral_7b", "llama3_8b"],
            domains=[
                "glue_mnli",   # ~285K - 多段落推理
                "glue_qnli",   # ~98K  - 问答推理
                "glue_sst2",   # ~24K  - 情感分析
                "glue_stsb",   # ~3.5K - 语义相似度
                "glue_cola",   # ~2K   - 语法判断
            ],
            limit_per_domain=limit_per_domain,
            output_dir="experiments/output/glue",
        )

    @classmethod
    def full_config(cls) -> "ExperimentConfig":
        """完整实验配置"""
        return cls(
            model_names=[
                "qwen_7b", "mistral_7b",
                "gemma_9b", "deepseek_16b", "internlm_7b"
            ],
            domains=[
                "news_en", "news_zh", "academic",
                "code", "dialogue", "literature"
            ],
            limit_per_domain=None,
            output_dir="experiments/output/full",
        )


@dataclass
class ModelSpectrumResult:
    """单个模型的谱计算结果"""
    model_name: str
    domain: str
    r_eff: float
    prob: np.ndarray
    eigvals: np.ndarray
    d: int  # 特征维度
    N: int  # 样本数


@dataclass
class ExperimentResult:
    """完整实验结果"""
    config: ExperimentConfig
    timestamp: str
    results: Dict[str, ModelSpectrumResult]  # key = "{model}__{domain}"

    # 距离矩阵
    distance_matrices: Dict[str, np.ndarray]  # key = domain, value = n×n W1 距离矩阵
    model_names: List[str]

    # 汇总统计
    summary: Dict[str, Any] = field(default_factory=dict)

    def save(self, output_dir: str) -> str:
        """保存实验结果"""
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        # 保存配置
        with open(output_path / "config.json", "w") as f:
            json.dump({
                "model_names": self.config.model_names,
                "domains": self.config.domains,
                "limit_per_domain": self.config.limit_per_domain,
                "normalize_method": self.config.normalize_method,
            }, f, indent=2)

        # 保存谱结果
        spectrum_results = {}
        for key, result in self.results.items():
            spectrum_results[key] = {
                "model_name": result.model_name,
                "domain": result.domain,
                "r_eff": result.r_eff,
                "d": result.d,
                "N": result.N,
                "prob": result.prob.tolist(),
                "eigvals": result.eigvals.tolist(),
            }
        with open(output_path / "spectrum_results.json", "w") as f:
            json.dump(spectrum_results, f, indent=2)

        # 保存距离矩阵
        dist_dir = output_path / "distance_matrices"
        dist_dir.mkdir(exist_ok=True)
        for domain, matrix in self.distance_matrices.items():
            np.save(dist_dir / f"{domain}_w1_matrix.npy", matrix)

        # 保存汇总统计
        with open(output_path / "summary.json", "w") as f:
            json.dump(self.summary, f, indent=2, default=str)

        # 保存元信息
        with open(output_path / "metadata.json", "w") as f:
            json.dump({
                "timestamp": self.timestamp,
                "num_models": len(self.model_names),
                "num_domains": len(self.config.domains),
            }, f, indent=2)

        return str(output_path)


# 导入 numpy（在 dataclass 之后）
import numpy as np


class VerificationExperiment:
    """
    LLM 有效秩谱收敛验证实验
    """

    def __init__(self, config: ExperimentConfig):
        self.config = config
        self.rep_loader = RepresentationLoader("model/representations")
        self.corpus_loader = CorpusLoader("database/corpus")

        # 结果存储
        self.results: Dict[str, ModelSpectrumResult] = {}
        self.distance_matrices: Dict[str, np.ndarray] = {}

    def run(self) -> ExperimentResult:
        """运行完整实验流程"""
        timestamp = datetime.now().isoformat()

        print(f"=== LLM 有效秩谱收敛验证实验 ===")
        print(f"时间：{timestamp}")
        print(f"模型数：{len(self.config.model_names)}")
        print(f"领域数：{len(self.config.domains)}")
        print(f"归一化方法：{self.config.normalize_method}")
        print()

        # Step 1: 计算每个模型的谱
        self._compute_spectra()

        # Step 2: 计算每个领域内的跨模型谱距离矩阵
        self._compute_distance_matrices()

        # Step 3: 生成汇总统计
        summary = self._generate_summary()

        return ExperimentResult(
            config=self.config,
            timestamp=timestamp,
            results=self.results,
            distance_matrices=self.distance_matrices,
            model_names=self.config.model_names,
            summary=summary,
        )

    def _compute_spectra(self) -> None:
        """Step 1: 计算每个模型的谱有效秩"""
        print("Step 1: 计算谱有效秩...")

        for model_name in self.config.model_names:
            for domain in self.config.domains:
                key = f"{model_name}__{domain}"
                print(f"  处理：{key}")

                # 加载表示矩阵
                H = self.rep_loader.load_representation(model_name, domain)

                # 计算谱
                r_eff, prob = compute_gram_spectrum(
                    H,
                    normalize=self.config.normalize_method,
                    return_eigvals=True
                )

                # 计算完整谱信息
                eigvals, _ = compute_spectrum_array(H, normalize=self.config.normalize_method)

                # 存储结果
                self.results[key] = ModelSpectrumResult(
                    model_name=model_name,
                    domain=domain,
                    r_eff=r_eff,
                    prob=prob,
                    eigvals=eigvals,
                    d=H.shape[0],
                    N=H.shape[1],
                )

                print(f"    r_eff = {r_eff:.4f}, d = {H.shape[0]}, N = {H.shape[1]}")

        print()

    def _compute_distance_matrices(self) -> None:
        """Step 2: 计算每个领域内的跨模型谱距离矩阵"""
        print("Step 2: 计算谱距离矩阵...")

        for domain in self.config.domains:
            print(f"  领域：{domain}")

            # 提取该领域所有模型的谱分布
            prob_dict = {}
            for model_name in self.config.model_names:
                key = f"{model_name}__{domain}"
                if key in self.results:
                    prob_dict[model_name] = self.results[key].prob

            # 计算距离矩阵
            if len(prob_dict) >= 2:
                dist_matrix, model_names = compute_spectrum_distance_matrix(
                    prob_dict, distance_fn="wasserstein1"
                )
                self.distance_matrices[domain] = dist_matrix

                print(f"    距离矩阵形状：{dist_matrix.shape}")
                print(f"    平均距离：{dist_matrix[np.triu_indices(len(model_names), k=1)].mean():.4f}")
            else:
                print(f"    跳过（模型数不足）")

        print()

    def _generate_summary(self) -> Dict[str, Any]:
        """Step 3: 生成汇总统计"""
        print("Step 3: 生成汇总统计...")

        summary = {
            "spectrum_statistics": {},
            "convergence_analysis": {},
            "per_domain": {},
        }

        # 谱统计
        for model_name in self.config.model_names:
            model_results = [
                r for key, r in self.results.items()
                if r.model_name == model_name
            ]
            if model_results:
                r_effs = [r.r_eff for r in model_results]
                summary["spectrum_statistics"][model_name] = {
                    "r_eff_mean": float(np.mean(r_effs)),
                    "r_eff_std": float(np.std(r_effs)),
                    "r_eff_min": float(np.min(r_effs)),
                    "r_eff_max": float(np.max(r_effs)),
                }

        # 跨模型收敛分析
        for domain, dist_matrix in self.distance_matrices.items():
            n = dist_matrix.shape[0]
            if n >= 2:
                upper_tri = dist_matrix[np.triu_indices(n, k=1)]
                summary["per_domain"][domain] = {
                    "mean_w1_distance": float(upper_tri.mean()),
                    "std_w1_distance": float(upper_tri.std()),
                    "max_w1_distance": float(upper_tri.max()),
                    "min_w1_distance": float(upper_tri.min()),
                }

        # 整体收敛判断
        all_distances = []
        for dist_matrix in self.distance_matrices.values():
            n = dist_matrix.shape[0]
            if n >= 2:
                all_distances.extend(dist_matrix[np.triu_indices(n, k=1)].tolist())

        if all_distances:
            overall_mean = float(np.mean(all_distances))
            summary["convergence_analysis"] = {
                "overall_mean_w1_distance": overall_mean,
                "convergence_judgment": "收敛" if overall_mean < 0.05 else "部分收敛" if overall_mean < 0.2 else "发散",
                "num_domain_matrices": len(self.distance_matrices),
            }

        print(f"  整体平均 W1 距离：{summary['convergence_analysis'].get('overall_mean_w1_distance', 'N/A')}")
        print(f"  收敛判断：{summary['convergence_analysis'].get('convergence_judgment', 'N/A')}")
        print()

        return summary


def run_mvp_experiment(
    output_dir: Optional[str] = None,
) -> ExperimentResult:
    """
    运行 MVP 实验（最小可行实验）

    配置：
        - 2 个模型：qwen_7b, mistral_7b
        - 6 个领域：news_en, news_zh, academic, code, dialogue, literature
        - 每个领域 500 条文本
    """
    config = ExperimentConfig.mvp_config()
    if output_dir:
        config.output_dir = output_dir

    experiment = VerificationExperiment(config)
    result = experiment.run()

    # 保存结果
    output_path = result.save(config.output_dir)
    print(f"结果已保存到：{output_path}")

    return result


def run_full_experiment(
    output_dir: Optional[str] = None,
) -> ExperimentResult:
    """
    运行完整实验

    配置：
        - 7 个模型
        - 6 个领域
        - 无文本数量限制
    """
    config = ExperimentConfig.full_config()
    if output_dir:
        config.output_dir = output_dir

    experiment = VerificationExperiment(config)
    result = experiment.run()

    # 保存结果
    output_path = result.save(config.output_dir)
    print(f"结果已保存到：{output_path}")

    return result


def run_scale_analysis(
    model_family: str = "llama",
    output_dir: str = "experiments/output/scale_analysis",
) -> Dict[str, Any]:
    """
    运行规模效应分析

    分析同一家族不同规模模型的谱收敛趋势
    """
    # 需要预先定义模型家族配置
    scale_configs = {
        "llama": [
            "llama_1b", "llama_3b", "llama_8b", "llama_70b"
        ],
        "qwen": [
            "qwen_0.5b", "qwen_1.5b", "qwen_3b", "qwen_7b", "qwen_72b"
        ],
    }

    if model_family not in scale_configs:
        raise ValueError(f"Unknown model family: {model_family}")

    model_names = scale_configs[model_family]
    domains = ["news_en", "academic"]  # 简化版

    config = ExperimentConfig(
        model_names=model_names,
        domains=domains,
        output_dir=output_dir,
    )

    experiment = VerificationExperiment(config)
    result = experiment.run()

    # 分析规模与谱收敛的关系
    scale_analysis = {
        "model_family": model_family,
        "model_names": model_names,
        "r_eff_by_model": {},
        "cross_model_distances": {},
    }

    for model_name in model_names:
        model_results = [
            r for key, r in result.results.items()
            if r.model_name == model_name
        ]
        r_effs = [r.r_eff for r in model_results]
        scale_analysis["r_eff_by_model"][model_name] = {
            "mean": float(np.mean(r_effs)),
            "std": float(np.std(r_effs)),
        }

    for domain, dist_matrix in result.distance_matrices.items():
        scale_analysis["cross_model_distances"][domain] = {
            "matrix": dist_matrix.tolist(),
            "mean": float(dist_matrix[np.triu_indices(len(model_names), k=1)].mean()),
        }

    # 保存规模分析结果
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    with open(output_path / "scale_analysis.json", "w") as f:
        json.dump(scale_analysis, f, indent=2)

    return scale_analysis


def run_glue_experiment(
    output_dir: Optional[str] = None,
    limit_per_domain: Optional[int] = 5000,
) -> ExperimentResult:
    """
    运行 GLUE Benchmark 实验

    配置：
        - 3 个模型：qwen_7b, mistral_7b, llama3_8b
        - 5 个 GLUE 子集：MNLI, QNLI, SST-2, STS-B, CoLA
        - 每个子集限制 5000 条样本（默认）
    """
    config = ExperimentConfig.glue_config(limit_per_domain=limit_per_domain)
    if output_dir:
        config.output_dir = output_dir

    experiment = VerificationExperiment(config)
    result = experiment.run()

    # 保存结果
    output_path = result.save(config.output_dir)
    print(f"结果已保存到：{output_path}")

    return result


# ============= 命令行入口 =============

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="LLM 有效秩谱收敛验证实验"
    )
    parser.add_argument(
        "--mode",
        choices=["mvp", "full", "scale", "glue"],
        default="mvp",
        help="实验模式：mvp=最小可行实验，full=完整实验，scale=规模效应分析，glue=GLUE 基准实验"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="输出目录（可选）"
    )
    parser.add_argument(
        "--model-family",
        type=str,
        default="llama",
        help="规模效应分析的模型家族（仅 scale 模式）"
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=5000,
        help="GLUE 实验每个领域的样本数限制（仅 glue 模式）"
    )

    args = parser.parse_args()

    if args.mode == "mvp":
        run_mvp_experiment(args.output)
    elif args.mode == "full":
        run_full_experiment(args.output)
    elif args.mode == "scale":
        run_scale_analysis(args.model_family, args.output or "experiments/output/scale_analysis")