05 - 实际项目经验分享¶

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 1. 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 2. 配置INT4量化
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# 3. 加载量化模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=quantization_config,
    device_map="auto"
)

print(f"模型大小: {model.get_memory_footprint() / 1e9:.2f} GB")

from fastapi import FastAPI
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = FastAPI(title="多模型量化服务")

class ModelManager:
    """
    模型管理器
    """
    def __init__(self):
        self.models = {}
        self.tokenizers = {}
        self.current_model = None

    def load_model(self, model_name, quantization="int4"):
        """
        加载模型
        """
        if model_name in self.models:
            self.current_model = model_name
            return

        # 加载分词器
        tokenizer = AutoTokenizer.from_pretrained(model_name)

        # 加载量化模型
        if quantization == "int4":
            from transformers import BitsAndBytesConfig
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16
            )
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                quantization_config=quantization_config,
                device_map="auto"
            )

        self.models[model_name] = model
        self.tokenizers[model_name] = tokenizer
        self.current_model = model_name

    def switch_model(self, model_name):
        """
        切换模型
        """
        if model_name in self.models:
            self.current_model = model_name
        else:
            self.load_model(model_name)

# 全局模型管理器
model_manager = ModelManager()

# 预加载模型
model_manager.load_model("meta-llama/Llama-2-7b-hf")
model_manager.load_model("mistralai/Mistral-7B-v0.1")

@app.post("/generate")
async def generate(prompt: str, model_name: str = None):
    """
    生成文本
    """
    # 切换模型
    if model_name:
        model_manager.switch_model(model_name)

    # 获取当前模型
    model = model_manager.models[model_manager.current_model]
    tokenizer = model_manager.tokenizers[model_manager.current_model]

    # 生成文本
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_length=200)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return {"result": result}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

def progressive_quantization(model, dataloader):
    """
    渐进式量化

    先量化到INT8，微调后再量化到INT4
    """
    # 第一步：量化到INT8
    print("量化到INT8...")
    int8_model = quantize_to_int8(model, dataloader)

    # 第二步：微调INT8模型
    print("微调INT8模型...")
    finetuned_int8 = finetune(int8_model, dataloader, epochs=3)

    # 第三步：量化到INT4
    print("量化到INT4...")
    int4_model = quantize_to_int4(finetuned_int8, dataloader)

    # 第四步：微调INT4模型
    print("微调INT4模型...")
    final_model = finetune(int4_model, dataloader, epochs=2)

    return final_model

class DynamicBatchProcessor:
    """
    动态批处理器
    """
    def __init__(self, model, max_batch_size=8, max_wait_time=0.1):
        self.model = model
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time
        self.requests = []

    async def add_request(self, request):  # async def定义异步函数；用await调用
        """
        添加请求
        """
        self.requests.append(request)

        # 检查是否需要处理
        if len(self.requests) >= self.max_batch_size:
            await self.process_batch()  # await等待异步操作完成

    async def process_batch(self):
        """
        处理批次
        """
        if not self.requests:
            return

        # 处理请求
        batch = self.requests[:self.max_batch_size]  # 切片操作：[start:end:step]提取子序列
        results = await self.model.generate_batch(batch)

        # 返回结果
        for request, result in zip(batch, results):  # zip并行遍历多个可迭代对象
            request.set_result(result)

        # 清空已处理的请求
        self.requests = self.requests[len(batch):]

class SmartCache:
    """
    智能缓存
    """
    def __init__(self, max_size=1000):
        self.cache = {}
        self.max_size = max_size
        self.access_count = {}

    def get(self, key):
        """
        获取缓存
        """
        if key in self.cache:
            self.access_count[key] += 1
            return self.cache[key]
        return None

    def set(self, key, value, priority=1):
        """
        设置缓存
        """
        # 如果缓存已满，删除最少使用的条目
        if len(self.cache) >= self.max_size:
            lru_key = min(self.access_count.keys(),
                         key=lambda k: (self.access_count[k], priority))  # lambda匿名函数：简洁的单行函数
            del self.cache[lru_key]
            del self.access_count[lru_key]

        self.cache[key] = value
        self.access_count[key] = 1