# 抽象基类（ABC）用于定义模型加载器的统一接口
from abc import ABC, abstractmethod
# 类型提示相关导入
from typing import Dict, Tuple, Union, Optional

# PyTorch 核心库
import torch
# Hugging Face 相关组件
from transformers import PreTrainedModel, PreTrainedTokenizer, AutoProcessor, BitsAndBytesConfig

class BaseModelLoader(ABC):
    """
    模型加载器基类，定义模型加载的标准流程
    核心功能：
    1. 统一管理模型加载参数
    2. 支持量化配置
    3. 支持硬件加速选项
    4. 提供抽象接口强制子类实现
    """
    
    def __init__(
        self, 
        model_hf_path: str,                              # Hugging Face 仓库模型标识
        model_local_path: str,                           # 本地模型目录路径
        compute_dtype: torch.dtype,                      # 计算数据类型（如 torch.float16）
        bnb_config: Optional[BitsAndBytesConfig] = None, # 量化配置（4/8-bit）
        use_flash_attn: bool = False,                    # 是否启用Flash Attention加速
        device_map: Optional[Union[Dict, str]] = None,   # 设备映射策略
    ) -> None:
        # 初始化模型路径参数
        self.model_hf_path = model_hf_path
        self.model_local_path = model_local_path
        
        # 构建模型加载的关键参数字典（技术难点：动态参数组合）
        self.loading_kwargs = dict(
            torch_dtype=compute_dtype,         # 控制模型权重数据类型
            quantization_config=bnb_config,    # 量化配置（QLoRA等场景）
            device_map=device_map,             # 分布式设备映射策略
            # 示例: device_map="auto" 表示自动分配多GPU/NPU设备
        )
        
        # Flash Attention 加速配置（硬件和模型兼容性要求高）
        if use_flash_attn:
            # 需要安装 flash-attn 库且硬件支持
            self.loading_kwargs["attn_implementation"] = "flash_attention_2"
            """
            技术难点：
            1. 需要CUDA架构>= SM80（如A100）或特定NPU支持
            2. 与量化配置可能存在兼容性问题
            3. 注意力模式需要模型结构支持
            """

    @abstractmethod
    def load(self, load_model: bool = True) -> Tuple[
        PreTrainedModel, 
        Union[None, PreTrainedTokenizer], 
        Union[None, AutoProcessor]
    ]: 
        """
        抽象方法：强制子类实现模型加载逻辑
        
        参数说明：
        load_model - 是否实际加载模型权重（可用于仅加载分词器/处理器）
        
        返回值三元组：
        1. PreTrainedModel: 加载的模型实例
        2. PreTrainedTokenizer: 文本分词器（多模态模型可能为None）
        3. AutoProcessor: 多模态处理器（纯文本模型可能为None）
        
        设计难点：
        1. 需要处理多种模型类型（纯文本/多模态）
        2. 需要协调分词器与处理器的加载顺序
        3. 需要处理本地缓存与远程仓库的优先级
        """
        ...