maque 0.2.1__py3-none-any.whl

maque/quantization/__init__.py

@@ -0,0 +1,42 @@
+"""模型量化模块
+
+提供多种量化方案的统一接口，支持 vLLM 推理和 QLoRA 微调场景。
+
+支持的量化方法:
+    - auto-round: Intel SGD 优化权重舍入，精度好 (推荐)
+    - awq: Activation-aware Weight Quantization
+    - gptq: 经典 GPTQ 量化
+    - bnb-nf4: 4-bit NormalFloat 量化 (QLoRA)
+    - bnb-int8: 8-bit 整数量化
+
+Examples:
+    >>> from maque.quantization import get_quantizer
+    >>> quantizer = get_quantizer("auto-round")
+    >>> quantizer.quantize("Qwen/Qwen3-4B", "./Qwen3-4B-quant")
+
+    >>> from maque.quantization import AutoRoundQuantizer
+    >>> quantizer = AutoRoundQuantizer(bits=4, group_size=128)
+    >>> quantizer.quantize(model_path, output_path)
+"""
+
+from .base import (
+    BaseQuantizer,
+    QuantConfig,
+    get_quantizer,
+    list_methods,
+    QUANTIZATION_METHODS,
+)
+from .auto_round import AutoRoundQuantizer
+from .llm_compressor import LLMCompressorQuantizer
+from .bitsandbytes import BitsAndBytesQuantizer
+
+__all__ = [
+    "BaseQuantizer",
+    "QuantConfig",
+    "get_quantizer",
+    "list_methods",
+    "QUANTIZATION_METHODS",
+    "AutoRoundQuantizer",
+    "LLMCompressorQuantizer",
+    "BitsAndBytesQuantizer",
+]

maque/quantization/auto_round.py

@@ -0,0 +1,120 @@
+"""AutoRound 量化器
+
+使用 Intel 的 auto-round 库进行量化，采用 SGD 优化权重舍入，精度损失小。
+"""
+
+from .base import BaseQuantizer, QuantConfig
+from typing import Optional
+from pathlib import Path
+
+
+class AutoRoundQuantizer(BaseQuantizer):
+    """AutoRound 量化器
+
+    使用 Intel 的 auto-round 库，通过 SGD 优化权重舍入实现高质量量化。
+
+    Args:
+        bits: 量化位数，默认 4
+        group_size: 量化分组大小，默认 128
+        sym: 是否对称量化，默认 True
+        iters: 优化迭代次数，默认 200
+        seqlen: 校准序列长度，默认 512
+        nsamples: 校准样本数，默认 256
+        batch_size: 批次大小，默认 4
+        low_gpu_mem_usage: 低显存模式，默认 True
+        format: 输出格式 (auto_round, auto_gptq)，默认 auto_round
+        dataset: 校准数据集，默认 NeelNanda/pile-10k
+
+    Examples:
+        >>> from maque.quantization import AutoRoundQuantizer
+        >>> quantizer = AutoRoundQuantizer(bits=4)
+        >>> quantizer.quantize("Qwen/Qwen3-4B", "./Qwen3-4B-quant")
+
+        # 使用自定义数据集
+        >>> quantizer = AutoRoundQuantizer(dataset="wikitext2")
+        >>> quantizer.quantize(model_path, output_path)
+    """
+
+    def __init__(
+        self,
+        bits: int = 4,
+        group_size: int = 128,
+        sym: bool = True,
+        iters: int = 200,
+        seqlen: int = 512,
+        nsamples: int = 256,
+        batch_size: int = 4,
+        low_gpu_mem_usage: bool = True,
+        format: str = "auto_round",
+        dataset: str = "NeelNanda/pile-10k",
+        **kwargs,
+    ):
+        config = QuantConfig(
+            bits=bits,
+            group_size=group_size,
+            sym=sym,
+            seqlen=seqlen,
+            nsamples=nsamples,
+            batch_size=batch_size,
+            low_gpu_mem_usage=low_gpu_mem_usage,
+        )
+        super().__init__(config)
+        self.iters = iters
+        self.format = format
+        self.dataset = dataset
+
+    @property
+    def method_name(self) -> str:
+        return "auto-round"
+
+    @property
+    def supported_formats(self):
+        return ["auto_round", "auto_gptq"]
+
+    def quantize(self, model_path: str, output_path: str, **kwargs) -> str:
+        """执行量化
+
+        Args:
+            model_path: 原始模型路径
+            output_path: 量化后模型保存路径
+            **kwargs: 额外参数传递给 AutoRound
+
+        Returns:
+            量化后模型路径
+        """
+        try:
+            from auto_round import AutoRound
+        except ImportError:
+            raise ImportError(
+                "auto-round 未安装，请运行: pip install auto-round"
+            )
+
+        output_path = Path(output_path)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        print(f"[auto-round] 加载模型: {model_path}")
+        print(f"[auto-round] 配置: bits={self.config.bits}, group_size={self.config.group_size}, "
+              f"seqlen={self.config.seqlen}, nsamples={self.config.nsamples}")
+        print(f"[auto-round] 校准数据集: {self.dataset}")
+
+        # 创建 AutoRound 实例
+        autoround = AutoRound(
+            model=model_path,
+            scheme="W4A16" if self.config.bits == 4 else f"W{self.config.bits}A16",
+            iters=self.iters,
+            seqlen=self.config.seqlen,
+            nsamples=self.config.nsamples,
+            batch_size=self.config.batch_size,
+            low_gpu_mem_usage=self.config.low_gpu_mem_usage,
+            dataset=self.dataset,
+            **kwargs,
+        )
+
+        print(f"[auto-round] 开始量化 (iters={self.iters})...")
+        autoround.quantize()
+
+        print(f"[auto-round] 保存到: {output_path}")
+        autoround.save_quantized(str(output_path), format=self.format)
+
+        print(f"[auto-round] 量化完成!")
+        return str(output_path)

maque/quantization/base.py

@@ -0,0 +1,145 @@
+"""量化器抽象基类"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Optional, List, Literal
+from pathlib import Path
+
+
+@dataclass
+class QuantConfig:
+    """量化配置"""
+    bits: int = 4
+    group_size: int = 128
+    sym: bool = True
+    seqlen: int = 512
+    nsamples: int = 128
+    batch_size: int = 4
+    low_gpu_mem_usage: bool = True
+
+
+class BaseQuantizer(ABC):
+    """量化器抽象基类"""
+
+    def __init__(self, config: QuantConfig = None):
+        self.config = config or QuantConfig()
+
+    @abstractmethod
+    def quantize(self, model_path: str, output_path: str, **kwargs) -> str:
+        """量化模型
+
+        Args:
+            model_path: 原始模型路径
+            output_path: 量化后模型保存路径
+            **kwargs: 额外参数
+
+        Returns:
+            量化后模型路径
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def method_name(self) -> str:
+        """量化方法名称"""
+        pass
+
+    @property
+    def supported_formats(self) -> List[str]:
+        """支持的输出格式"""
+        return ["auto"]
+
+    def get_model_info(self, model_path: str) -> dict:
+        """获取模型的量化信息"""
+        import json
+        model_path = Path(model_path)
+
+        # 检查 quantization_config.json
+        quant_config_path = model_path / "quantization_config.json"
+        if quant_config_path.exists():
+            with open(quant_config_path, "r") as f:
+                return json.load(f)
+
+        # 检查 config.json 中的 quantization_config
+        config_path = model_path / "config.json"
+        if config_path.exists():
+            with open(config_path, "r") as f:
+                config = json.load(f)
+                if "quantization_config" in config:
+                    return config["quantization_config"]
+
+        return {}
+
+
+# 支持的量化方法
+QUANTIZATION_METHODS = {
+    "auto-round": "AutoRoundQuantizer",
+    "awq": "LLMCompressorQuantizer",
+    "gptq": "LLMCompressorQuantizer",
+    "bnb-nf4": "BitsAndBytesQuantizer",
+    "bnb-int8": "BitsAndBytesQuantizer",
+}
+
+
+def get_quantizer(method: str, **kwargs) -> BaseQuantizer:
+    """根据方法名获取量化器
+
+    Args:
+        method: 量化方法名称 (auto-round, awq, gptq, bnb-nf4, bnb-int8)
+        **kwargs: 传递给量化器的参数
+
+    Returns:
+        BaseQuantizer 实例
+    """
+    if method not in QUANTIZATION_METHODS:
+        available = ", ".join(QUANTIZATION_METHODS.keys())
+        raise ValueError(f"不支持的量化方法: {method}，可用方法: {available}")
+
+    if method == "auto-round":
+        from .auto_round import AutoRoundQuantizer
+        return AutoRoundQuantizer(**kwargs)
+    elif method in ("awq", "gptq"):
+        from .llm_compressor import LLMCompressorQuantizer
+        return LLMCompressorQuantizer(scheme=method, **kwargs)
+    elif method in ("bnb-nf4", "bnb-int8"):
+        from .bitsandbytes import BitsAndBytesQuantizer
+        bits = 4 if method == "bnb-nf4" else 8
+        return BitsAndBytesQuantizer(bits=bits, **kwargs)
+    else:
+        raise ValueError(f"未实现的量化方法: {method}")
+
+
+def list_methods() -> dict:
+    """列出所有支持的量化方法及其描述"""
+    return {
+        "auto-round": {
+            "library": "auto-round",
+            "precision": "W4A16",
+            "description": "Intel 出品，SGD 优化权重舍入，精度好",
+            "use_case": "vLLM 推理",
+        },
+        "awq": {
+            "library": "llm-compressor",
+            "precision": "W4A16",
+            "description": "Activation-aware Weight Quantization",
+            "use_case": "vLLM 推理",
+        },
+        "gptq": {
+            "library": "llm-compressor",
+            "precision": "W4A16",
+            "description": "经典 GPTQ 量化",
+            "use_case": "通用推理",
+        },
+        "bnb-nf4": {
+            "library": "bitsandbytes",
+            "precision": "NF4",
+            "description": "4-bit NormalFloat 量化",
+            "use_case": "QLoRA 微调",
+        },
+        "bnb-int8": {
+            "library": "bitsandbytes",
+            "precision": "INT8",
+            "description": "8-bit 整数量化",
+            "use_case": "显存节省",
+        },
+    }

maque/quantization/bitsandbytes.py

@@ -0,0 +1,127 @@
+"""BitsAndBytes 量化器
+
+使用 bitsandbytes 库进行 NF4/INT8 量化，主要用于 QLoRA 微调场景。
+注意：bitsandbytes 是推理时动态量化，不生成独立的量化模型文件。
+"""
+
+from .base import BaseQuantizer, QuantConfig
+from typing import Literal
+from pathlib import Path
+
+
+class BitsAndBytesQuantizer(BaseQuantizer):
+    """BitsAndBytes 量化器
+
+    使用 bitsandbytes 库进行 NF4/INT8 量化。
+
+    注意：bitsandbytes 是推理时动态量化，调用 quantize() 会：
+    1. 加载模型并应用量化配置
+    2. 保存带有 quantization_config 的模型配置
+
+    加载时需要使用 load_in_4bit=True 或 load_in_8bit=True。
+
+    Args:
+        bits: 量化位数 (4 或 8)，默认 4
+        bnb_4bit_compute_dtype: 4bit 计算精度，默认 bfloat16
+        bnb_4bit_quant_type: 4bit 量化类型 (nf4, fp4)，默认 nf4
+        bnb_4bit_use_double_quant: 是否使用双重量化，默认 True
+
+    Examples:
+        >>> from maque.quantization import BitsAndBytesQuantizer
+        >>> quantizer = BitsAndBytesQuantizer(bits=4)
+        >>> quantizer.quantize("Qwen/Qwen3-4B", "./Qwen3-4B-bnb")
+    """
+
+    def __init__(
+        self,
+        bits: Literal[4, 8] = 4,
+        bnb_4bit_compute_dtype: str = "bfloat16",
+        bnb_4bit_quant_type: Literal["nf4", "fp4"] = "nf4",
+        bnb_4bit_use_double_quant: bool = True,
+        **kwargs,
+    ):
+        config = QuantConfig(bits=bits)
+        super().__init__(config)
+        self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
+        self.bnb_4bit_quant_type = bnb_4bit_quant_type
+        self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
+
+    @property
+    def method_name(self) -> str:
+        return f"bnb-{'nf4' if self.config.bits == 4 else 'int8'}"
+
+    def quantize(self, model_path: str, output_path: str, **kwargs) -> str:
+        """应用量化配置并保存模型
+
+        注意：bitsandbytes 是推理时动态量化，此方法会：
+        1. 使用量化配置加载模型
+        2. 保存模型和带有 quantization_config 的配置文件
+
+        Args:
+            model_path: 原始模型路径
+            output_path: 输出路径
+            **kwargs: 额外参数
+
+        Returns:
+            输出路径
+        """
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+        except ImportError as e:
+            if "bitsandbytes" in str(e):
+                raise ImportError(
+                    "bitsandbytes 未安装，请运行: pip install bitsandbytes"
+                )
+            raise
+
+        output_path = Path(output_path)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        print(f"[bnb] 加载模型: {model_path}")
+
+        # 配置 BitsAndBytes
+        if self.config.bits == 4:
+            compute_dtype = getattr(torch, self.bnb_4bit_compute_dtype)
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=compute_dtype,
+                bnb_4bit_quant_type=self.bnb_4bit_quant_type,
+                bnb_4bit_use_double_quant=self.bnb_4bit_use_double_quant,
+            )
+            print(f"[bnb] 配置: NF4, compute_dtype={self.bnb_4bit_compute_dtype}")
+        else:
+            bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+            print(f"[bnb] 配置: INT8")
+
+        # 加载模型
+        print(f"[bnb] 应用量化配置加载模型...")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            quantization_config=bnb_config,
+            device_map="auto",
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+        # 保存模型和配置
+        print(f"[bnb] 保存到: {output_path}")
+        model.save_pretrained(str(output_path))
+        tokenizer.save_pretrained(str(output_path))
+
+        print(f"[bnb] 完成! 加载时请使用 load_in_{self.config.bits}bit=True")
+        return str(output_path)
+
+    def get_load_kwargs(self) -> dict:
+        """获取加载量化模型时需要的参数"""
+        import torch
+
+        if self.config.bits == 4:
+            compute_dtype = getattr(torch, self.bnb_4bit_compute_dtype)
+            return {
+                "load_in_4bit": True,
+                "bnb_4bit_compute_dtype": compute_dtype,
+                "bnb_4bit_quant_type": self.bnb_4bit_quant_type,
+                "bnb_4bit_use_double_quant": self.bnb_4bit_use_double_quant,
+            }
+        else:
+            return {"load_in_8bit": True}

maque/quantization/llm_compressor.py

@@ -0,0 +1,102 @@
+"""LLM Compressor 量化器
+
+使用 vLLM 官方的 llm-compressor 库进行 AWQ/GPTQ 量化。
+"""
+
+from .base import BaseQuantizer, QuantConfig
+from typing import Literal
+from pathlib import Path
+
+
+class LLMCompressorQuantizer(BaseQuantizer):
+    """LLM Compressor 量化器
+
+    使用 vLLM 官方的 llm-compressor 库，支持 AWQ 和 GPTQ 量化方案。
+
+    Args:
+        scheme: 量化方案 (awq, gptq)，默认 awq
+        bits: 量化位数，默认 4
+        group_size: 量化分组大小，默认 128
+        sym: 是否对称量化，默认 True
+
+    Examples:
+        >>> from maque.quantization import LLMCompressorQuantizer
+        >>> quantizer = LLMCompressorQuantizer(scheme="awq")
+        >>> quantizer.quantize("Qwen/Qwen3-4B", "./Qwen3-4B-awq")
+    """
+
+    def __init__(
+        self,
+        scheme: Literal["awq", "gptq"] = "awq",
+        bits: int = 4,
+        group_size: int = 128,
+        sym: bool = True,
+        **kwargs,
+    ):
+        config = QuantConfig(
+            bits=bits,
+            group_size=group_size,
+            sym=sym,
+        )
+        super().__init__(config)
+        self.scheme = scheme
+
+    @property
+    def method_name(self) -> str:
+        return self.scheme
+
+    def quantize(self, model_path: str, output_path: str, **kwargs) -> str:
+        """执行量化
+
+        Args:
+            model_path: 原始模型路径
+            output_path: 量化后模型保存路径
+            **kwargs: 额外参数
+
+        Returns:
+            量化后模型路径
+        """
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+            from llmcompressor.modifiers.quantization import QuantizationModifier
+            from llmcompressor import oneshot
+        except ImportError as e:
+            if "llmcompressor" in str(e) or "oneshot" in str(e):
+                raise ImportError(
+                    "llm-compressor 未安装，请运行: pip install llmcompressor"
+                )
+            raise
+
+        output_path = Path(output_path)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        print(f"[{self.scheme}] 加载模型: {model_path}")
+
+        # 加载模型
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            device_map="auto",
+            torch_dtype="auto",
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+        # 配置量化方案
+        scheme_name = f"W{self.config.bits}A16"
+        print(f"[{self.scheme}] 配置: scheme={scheme_name}, group_size={self.config.group_size}")
+
+        recipe = QuantizationModifier(
+            targets="Linear",
+            scheme=scheme_name,
+            ignore=["lm_head"],
+        )
+
+        print(f"[{self.scheme}] 开始量化...")
+        oneshot(
+            model=model,
+            tokenizer=tokenizer,
+            recipe=recipe,
+            output_dir=str(output_path),
+        )
+
+        print(f"[{self.scheme}] 量化完成! 保存到: {output_path}")
+        return str(output_path)

maque/retriever/__init__.py

@@ -0,0 +1,35 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+Retriever 模块 - 提供向量检索功能
+
+支持 ChromaDB 和 Milvus 两种向量数据库后端，可独立使用。
+"""
+
+from .document import Document, SearchResult, Modality
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .chroma import ChromaRetriever
+    from .milvus import MilvusRetriever
+
+
+def __getattr__(name: str):
+    """延迟导入，避免未使用的依赖"""
+    if name == "ChromaRetriever":
+        from .chroma import ChromaRetriever
+        return ChromaRetriever
+    elif name == "MilvusRetriever":
+        from .milvus import MilvusRetriever
+        return MilvusRetriever
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "Document",
+    "SearchResult",
+    "Modality",
+    "ChromaRetriever",
+    "MilvusRetriever",
+]