evaluation/quantization_tool.py

"""
模型量化工具

完整的模型量化流程：
1. 动态量化（DyQuant）
2. 量化感知训练（QAT）
3. 量化后评估
4. 性能对比（量化 vs 原始）
"""

import sys
import time
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
from torch.quantization import get_default_qconfig, prepare_qat, convert

sys.path.insert(0, '.')

from evaluation.metrics import ModelEvaluator, EvaluationMetrics
from inference.dyquant import DyQuantConverter, QATTrainer, QuantConfig


class QuantizationTool:
    """模型量化工具（完整流程）"""
    
    def __init__(
        self,
        model: nn.Module,
        tokenizer = None,
        device: str = "cpu",
        model_path: Optional[str] = None,
    ):
        """
        初始化量化工具
        
        参数：
            model: 要量化的模型
            tokenizer: tokenizer（可选）
            device: 设备
            model_path: 模型路径（用于 DyQuant）
        """
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.model_path = model_path or "fusion-mini"
        self.original_model = None
        self.quantized_model = None
        self.qat_trainer = None
        self.converter = None
    
    def backup_original_model(self):
        """备份原始模型"""
        print("[QuantTool] 备份原始模型...")
        import copy
        self.original_model = copy.deepcopy(self.model)
        print("[QuantTool] 备份完成")
    
    def dynamic_quantize(
        self,
        bits: int = 8,
        mode: str = "symmetric",
    ) -> nn.Module:
        """
        动态量化（post-training quantization）
        
        参数：
            bits: 量化位数（4/8）
            mode: 量化模式（symmetric/asymmetric）
            
        返回：
            量化后的模型
        """
        print(f"[QuantTool] 开始动态量化（{bits}-bit, {mode}）...")
        
        # 使用正确的 QuantConfig API
        config = QuantConfig(
            model_path=self.model_path,
            bits=bits,
            mixed_precision=False,
        )
        
        self.converter = DyQuantConverter(config)
        self.converter.model = self.model  # 注入已加载的模型
        self.quantized_model = self.converter.convert()
        
        print(f"[QuantTool] 动态量化完成")
        return self.quantized_model
    
    def prepare_qat(
        self,
        learning_rate: float = 1e-4,
        num_epochs: int = 3,
        train_data: Optional[str] = None,
    ) -> "QATTrainer":
        """
        准备量化感知训练（QAT）
        
        参数：
            learning_rate: 学习率
            num_epochs: 训练轮数（注意：QATTrainer 没有 num_epochs 参数）
            train_data: 训练数据路径
            
        返回：
            QATTrainer 对象
        """
        print(f"[QuantTool] 准备量化感知训练（QAT）...")
        print(f"   学习率: {learning_rate}")
        print(f"   训练轮数: {num_epochs}")
        
        # 使用正确的 QuantConfig API
        config = QuantConfig(
            model_path=self.model_path,
            bits=4,
            mixed_precision=True,
        )
        
        # QATTrainer 签名：(config, train_data, learning_rate, warmup_steps)
        self.qat_trainer = QATTrainer(
            config=config,
            train_data=train_data,
            learning_rate=learning_rate,
            warmup_steps=100,
        )
        
        # 注入已加载的模型，避免重复加载
        self.qat_trainer.model = self.model
        
        self.qat_trainer.prepare()
        print(f"[QuantTool] QAT 准备完成")
        
        # 保存 num_epochs 供后续 train() 使用
        self._qat_epochs = num_epochs
        
        return self.qat_trainer
    
    def run_qat_training(self, epochs: Optional[int] = None):
        """运行 QAT 训练"""
        if self.qat_trainer is None:
            raise ValueError("请先调用 prepare_qat()")
        
        epochs = epochs or getattr(self, '_qat_epochs', 3)
        # QATTrainer.train() 签名：(epochs, lr, batch_size, max_len)
        self.qat_trainer.train(epochs=epochs)
        self.quantized_model = self.qat_trainer.qat_model
    
    def evaluate_quantized(
        self,
        texts: List[str],
        max_length: int = 512,
    ) -> EvaluationMetrics:
        """
        评估量化后的模型
        
        参数：
            texts: 评估文本
            max_length: 最大长度
            
        返回：
            评估指标
        """
        if self.quantized_model is None:
            raise ValueError("请先量化模型")
        
        print(f"[QuantTool] 评估量化后的模型...")
        
        evaluator = ModelEvaluator(
            model=self.quantized_model,
            tokenizer=self.tokenizer,
            device=self.device,
        )
        
        metrics = evaluator.evaluate(texts, max_length)
        print(f"[QuantTool] 评估完成")
        print(metrics)
        
        return metrics
    
    def compare_models(
        self,
        texts: List[str],
        max_length: int = 512,
    ) -> Dict[str, EvaluationMetrics]:
        """
        对比原始模型和量化模型的性能
        
        参数：
            texts: 评估文本
            max_length: 最大长度
            
        返回：
            包含原始模型和量化模型评估指标的字典
        """
        if self.original_model is None:
            self.backup_original_model()
        
        if self.quantized_model is None:
            raise ValueError("请先量化模型")
        
        print(f"[QuantTool] 对比原始模型和量化模型...")
        print()
        
        # 1. 评估原始模型
        print("[1] 评估原始模型...")
        original_evaluator = ModelEvaluator(
            model=self.original_model,
            tokenizer=self.tokenizer,
            device=self.device,
        )
        original_metrics = original_evaluator.evaluate(texts, max_length)
        print(f"   原始模型 Perplexity: {original_metrics.perplexity:.4f}")
        print(f"   原始模型 Loss: {original_metrics.loss:.4f}")
        print()
        
        # 2. 评估量化模型
        print("[2] 评估量化模型...")
        quantized_evaluator = ModelEvaluator(
            model=self.quantized_model,
            tokenizer=self.tokenizer,
            device=self.device,
        )
        quantized_metrics = quantized_evaluator.evaluate(texts, max_length)
        print(f"   量化模型 Perplexity: {quantized_metrics.perplexity:.4f}")
        print(f"   量化模型 Loss: {quantized_metrics.loss:.4f}")
        print()
        
        # 3. 计算差异
        print("[3] 性能差异...")
        perplexity_diff = quantized_metrics.perplexity - original_metrics.perplexity
        loss_diff = quantized_metrics.loss - original_metrics.loss
        
        print(f"   Perplexity 差异: {perplexity_diff:+.4f}")
        print(f"   Loss 差异: {loss_diff:+.4f}")
        print()
        
        # 4. 模型大小对比
        original_size = self._get_model_size(self.original_model)
        quantized_size = self._get_model_size(self.quantized_model)
        size_reduction = (1 - quantized_size / original_size) * 100
        
        print("[4] 模型大小对比...")
        print(f"   原始模型大小: {original_size / 1e6:.2f} MB")
        print(f"   量化模型大小: {quantized_size / 1e6:.2f} MB")
        print(f"   大小减少: {size_reduction:.2f}%")
        print()
        
        print("[QuantTool] 对比完成")
        
        return {
            "original": original_metrics,
            "quantized": quantized_metrics,
        }
    
    def _get_model_size(self, model: nn.Module) -> int:
        """获取模型大小（字节）"""
        total_size = 0
        for param in model.parameters():
            total_size += param.nelement() * param.element_size()
        return total_size
    
    def save_quantized(
        self,
        path: str,
        format: str = "safetensors",
    ):
        """
        保存量化后的模型
        
        参数：
            path: 保存路径
            format: 格式（safetensors/pytorch）
        """
        if self.quantized_model is None:
            raise ValueError("请先量化模型")
        
        print(f"[QuantTool] 保存量化模型到 {path}...")
        
        if format == "safetensors":
            try:
                import safetensors.torch
                safetensors.torch.save_file(self.quantized_model.state_dict(), path)
            except ImportError:
                print(f"[QuantTool] 警告：safetensors 未安装，使用 PyTorch 格式")
                format = "pytorch"
        
        if format == "pytorch":
            torch.save(self.quantized_model.state_dict(), path)
        
        print(f"[QuantTool] 保存完成")


def quantize_model(
    model: nn.Module,
    method: str = "dynamic",
    bits: int = 8,
    **kwargs,
) -> nn.Module:
    """
    便捷函数：量化模型
    
    参数：
        model: 要量化的模型
        method: 量化方法（dynamic/qat）
        bits: 量化位数
        **kwargs: 其他参数
        
    返回：
        量化后的模型
    """
    tool = QuantizationTool(model, **kwargs)
    
    if method == "dynamic":
        return tool.dynamic_quantize(bits=bits)
    elif method == "qat":
        trainer = tool.prepare_qat(**kwargs)
        epochs = kwargs.get('num_epochs', 3)
        trainer.train(epochs=epochs)
        return tool.quantized_model
    else:
        raise ValueError(f"不支持的量化方法: {method}")


if __name__ == "__main__":
    print("[QuantTool] 模型量化工具")
    print()
    print("功能：")
    print("  1. 动态量化（post-training）")
    print("  2. 量化感知训练（QAT）")
    print("  3. 量化后评估")
    print("  4. 性能对比（量化 vs 原始）")
    print()
    print("用法：")
    print("  from evaluation.quantization_tool import QuantizationTool")
    print("  tool = QuantizationTool(model, model_path='fusion-mini')")
    print("  quantized_model = tool.dynamic_quantize(bits=8)")
    print("  metrics = tool.compare_models(texts)")