06 - 实战项目：智能对话助手¶

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

class ModelManager:
    """
    模型管理器
    """
    def __init__(self, model_name="meta-llama/Llama-2-7b-hf"):
        self.model_name = model_name
        self.tokenizer = None
        self.model = None
        self.quantization_config = None

    def load_model(self, quantization="int4"):
        """
        加载模型

        Args:
            quantization: 量化类型 ("int4", "int8", "fp16", "fp32")
        """
        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)

        # 配置量化
        if quantization == "int4":
            self.quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4"
            )
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                quantization_config=self.quantization_config,
                device_map="auto"
            )
        elif quantization == "int8":
            self.quantization_config = BitsAndBytesConfig(
                load_in_8bit=True
            )
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                quantization_config=self.quantization_config,
                device_map="auto"
            )
        elif quantization == "fp16":
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=torch.float16,
                device_map="auto"
            )
        else:
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                device_map="auto"
            )

        print(f"模型加载完成！量化类型: {quantization}")
        print(f"模型大小: {self.model.get_memory_footprint() / 1e9:.2f} GB")

        return self.model, self.tokenizer

    def get_model_info(self):
        """
        获取模型信息
        """
        return {
            "model_name": self.model_name,
            "parameters": self.model.num_parameters(),
            "memory_footprint": self.model.get_memory_footprint() / 1e9,
            "device": str(self.model.device)
        }

# 使用示例
# model_manager = ModelManager("meta-llama/Llama-2-7b-hf")
# model, tokenizer = model_manager.load_model(quantization="int4")
# print(model_manager.get_model_info())

from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
import torch
import asyncio
from datetime import datetime
import hashlib
import json

app = FastAPI(title="智能对话助手API")

# 全局变量
model_manager = None
cache = {}

class ChatRequest(BaseModel):  # BaseModel Pydantic数据验证模型
    message: str
    conversation_id: str | None = None
    max_tokens: int = 200
    temperature: float = 0.7
    top_p: float = 0.95
    use_cache: bool = True

class ChatResponse(BaseModel):
    response: str
    conversation_id: str
    tokens_generated: int
    from_cache: bool

class ConversationHistory:
    """
    对话历史管理
    """
    def __init__(self):
        self.conversations = {}

    def add_message(self, conversation_id: str, role: str, content: str):
        """
        添加消息到对话历史
        """
        if conversation_id not in self.conversations:
            self.conversations[conversation_id] = []

        self.conversations[conversation_id].append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })

    def get_history(self, conversation_id: str) -> str:
        """
        获取对话历史
        """
        if conversation_id not in self.conversations:
            return ""

        history = self.conversations[conversation_id]
        history_text = "\n".join([
            f"{msg['role']}: {msg['content']}"
            for msg in history
        ])

        return history_text

# 全局对话历史
conversation_history = ConversationHistory()

def generate_cache_key(message: str, **kwargs) -> str:  # *args接收任意位置参数，**kwargs接收任意关键字参数
    """
    生成缓存键
    """
    data = {"message": message, **kwargs}
    data_str = json.dumps(data, sort_keys=True)  # json.dumps将Python对象序列化为JSON字符串
    return hashlib.md5(data_str.encode()).hexdigest()

@app.on_event("startup")
async def startup_event():  # async定义异步函数
    """
    启动事件
    """
    global model_manager
    model_manager = ModelManager("meta-llama/Llama-2-7b-hf")
    model_manager.load_model(quantization="int4")
    print("模型加载完成！")

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """
    对话接口
    """
    try:  # try/except捕获异常
        # 生成对话ID
        if not request.conversation_id:
            request.conversation_id = f"conv_{datetime.now().timestamp()}"

        # 检查缓存
        if request.use_cache:
            cache_key = generate_cache_key(
                request.message,
                max_tokens=request.max_tokens,
                temperature=request.temperature
            )

            if cache_key in cache:
                cached_response = cache[cache_key]
                conversation_history.add_message(
                    request.conversation_id,
                    "user",
                    request.message
                )
                conversation_history.add_message(
                    request.conversation_id,
                    "assistant",
                    cached_response
                )

                return ChatResponse(
                    response=cached_response,
                    conversation_id=request.conversation_id,
                    tokens_generated=len(cached_response.split()),
                    from_cache=True
                )

        # 获取对话历史
        history = conversation_history.get_history(request.conversation_id)
        prompt = f"{history}\nuser: {request.message}\nassistant:"

        # 编码输入
        inputs = model_manager.tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=2048
        ).to(model_manager.model.device)

        # 生成响应
        with torch.no_grad():  # 禁用梯度计算，节省内存
            outputs = model_manager.model.generate(
                **inputs,
                max_new_tokens=request.max_tokens,
                temperature=request.temperature,
                top_p=request.top_p,
                do_sample=True,
                pad_token_id=model_manager.tokenizer.eos_token_id
            )

        # 解码响应
        response = model_manager.tokenizer.decode(
            outputs[0],
            skip_special_tokens=True
        )

        # 提取assistant的响应
        if "assistant:" in response:
            response = response.split("assistant:")[-1].strip()  # [-1]负索引取最后元素

        # 更新对话历史
        conversation_history.add_message(
            request.conversation_id,
            "user",
            request.message
        )
        conversation_history.add_message(
            request.conversation_id,
            "assistant",
            response
        )

        # 缓存响应
        if request.use_cache:
            cache[cache_key] = response

        # 计算生成的token数
        # 使用更健壮的写法处理可能的1D/2D张量
        output_len = outputs.shape[-1] if outputs.dim() > 1 else outputs.shape[0]
        input_len = inputs["input_ids"].shape[-1] if inputs["input_ids"].dim() > 1 else inputs["input_ids"].shape[0]
        tokens_generated = output_len - input_len

        return ChatResponse(
            response=response,
            conversation_id=request.conversation_id,
            tokens_generated=tokens_generated,
            from_cache=False
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/conversations/{conversation_id}")
async def get_conversation(conversation_id: str):
    """
    获取对话历史
    """
    if conversation_id not in conversation_history.conversations:
        raise HTTPException(status_code=404, detail="Conversation not found")

    return {
        "conversation_id": conversation_id,
        "messages": conversation_history.conversations[conversation_id]
    }

@app.delete("/conversations/{conversation_id}")
async def delete_conversation(conversation_id: str):
    """
    删除对话
    """
    if conversation_id not in conversation_history.conversations:
        raise HTTPException(status_code=404, detail="Conversation not found")

    del conversation_history.conversations[conversation_id]

    return {"message": "Conversation deleted"}

@app.get("/health")
async def health_check():
    """
    健康检查
    """
    return {
        "status": "healthy",
        "model_loaded": model_manager is not None,
        "cache_size": len(cache)
    }

@app.get("/stats")
async def get_stats():
    """
    获取统计信息
    """
    return {
        "total_conversations": len(conversation_history.conversations),
        "cache_size": len(cache),
        "model_info": model_manager.get_model_info() if model_manager else None
    }

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

from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
import torch
import asyncio
from collections import defaultdict
import time

app = FastAPI(title="批处理推理服务")

# 全局变量
model_manager = None
batch_queue = defaultdict(list)  # defaultdict访问不存在的键时返回默认值
batch_results = {}
processing = False

class BatchRequest(BaseModel):
    messages: list[str]
    max_tokens: int = 200
    temperature: float = 0.7
    top_p: float = 0.95

class BatchResponse(BaseModel):
    request_id: str
    responses: list[str]
    processing_time: float

class BatchProcessor:
    """
    批处理器
    """
    def __init__(self, model_manager, max_batch_size=8, max_wait_time=0.5):
        self.model_manager = model_manager
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time
        self.processing = False

    async def add_request(self, request_id: str, messages: list[str], **kwargs):
        """
        添加请求到批次
        """
        batch_queue[request_id] = {
            "messages": messages,
            "kwargs": kwargs,
            "timestamp": time.time()
        }

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

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

        self.processing = True

        try:
            # 准备批次
            request_ids = list(batch_queue.keys())
            all_messages = []
            kwargs = None

            for request_id in request_ids:
                request_data = batch_queue[request_id]
                all_messages.extend(request_data["messages"])
                if kwargs is None:
                    kwargs = request_data["kwargs"]

            # 编码所有消息
            inputs = self.model_manager.tokenizer(
                all_messages,
                padding=True,
                truncation=True,
                return_tensors="pt"
            ).to(self.model_manager.model.device)

            # 批量生成
            start_time = time.time()

            with torch.no_grad():
                outputs = self.model_manager.model.generate(
                    **inputs,
                    max_new_tokens=kwargs.get("max_tokens", 200),
                    temperature=kwargs.get("temperature", 0.7),
                    top_p=kwargs.get("top_p", 0.95),
                    do_sample=True,
                    pad_token_id=self.model_manager.tokenizer.eos_token_id
                )

            processing_time = time.time() - start_time

            # 解码响应
            responses = []
            for output in outputs:
                response = self.model_manager.tokenizer.decode(
                    output,
                    skip_special_tokens=True
                )
                responses.append(response)

            # 分发结果
            msg_idx = 0
            for request_id in request_ids:
                request_data = batch_queue[request_id]
                num_messages = len(request_data["messages"])
                request_responses = responses[msg_idx:msg_idx + num_messages]

                batch_results[request_id] = {
                    "responses": request_responses,
                    "processing_time": processing_time
                }

                msg_idx += num_messages

            # 清空批次队列
            batch_queue.clear()

        finally:
            self.processing = False

# 全局批处理器
batch_processor = None

@app.on_event("startup")
async def startup_event():
    """
    启动事件
    """
    global model_manager, batch_processor

    model_manager = ModelManager("meta-llama/Llama-2-7b-hf")
    model_manager.load_model(quantization="int4")

    batch_processor = BatchProcessor(
        model_manager,
        max_batch_size=8,
        max_wait_time=0.5
    )

    print("批处理服务启动完成！")

@app.post("/batch-generate")
async def batch_generate(request: BatchRequest, background_tasks: BackgroundTasks):
    """
    批量生成
    """
    request_id = f"req_{time.time()}"

    # 添加请求到批次
    await batch_processor.add_request(
        request_id,
        request.messages,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p
    )

    # 返回请求ID
    return {"request_id": request_id}

@app.get("/batch-result/{request_id}")
async def get_batch_result(request_id: str):
    """
    获取批处理结果
    """
    if request_id not in batch_results:
        return {"status": "processing"}

    result = batch_results[request_id]
    del batch_results[request_id]

    return BatchResponse(
        request_id=request_id,
        responses=result["responses"],
        processing_time=result["processing_time"]
    )

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

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import redis
import json
from datetime import datetime, timedelta
import hashlib

app = FastAPI(title="缓存服务")

# Redis客户端
redis_client = redis.Redis(host='localhost', port=6379, db=0)

class CacheRequest(BaseModel):
    key: str
    value: str
    ttl: int | None = 3600

class CacheResponse(BaseModel):
    key: str
    value: str | None
    found: bool

def generate_cache_key(data: dict) -> str:
    """
    生成缓存键
    """
    data_str = json.dumps(data, sort_keys=True)
    return hashlib.md5(data_str.encode()).hexdigest()

@app.post("/cache/set")
async def set_cache(request: CacheRequest):
    """
    设置缓存
    """
    try:
        redis_client.setex(
            request.key,
            request.ttl,
            request.value
        )

        return {"message": "Cache set successfully", "key": request.key}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/cache/get/{key}", response_model=CacheResponse)
async def get_cache(key: str):
    """
    获取缓存
    """
    try:
        value = redis_client.get(key)

        if value is None:
            return CacheResponse(key=key, value=None, found=False)

        return CacheResponse(
            key=key,
            value=value.decode('utf-8'),
            found=True
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/cache/delete/{key}")
async def delete_cache(key: str):
    """
    删除缓存
    """
    try:
        redis_client.delete(key)
        return {"message": "Cache deleted successfully", "key": key}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/cache/stats")
async def get_cache_stats():
    """
    获取缓存统计
    """
    try:
        info = redis_client.info('stats')

        return {
            "total_keys": redis_client.dbsize(),
            "hits": info.get('keyspace_hits', 0),
            "misses": info.get('keyspace_misses', 0),
            "hit_rate": info.get('keyspace_hits', 0) / (
                info.get('keyspace_hits', 0) + info.get('keyspace_misses', 1)
            )
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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

from fastapi import FastAPI
from prometheus_client import Counter, Histogram, Gauge, generate_latest
from prometheus_client.exposition import CONTENT_TYPE_LATEST
from fastapi.responses import Response
import psutil
import torch

app = FastAPI(title="监控服务")

# Prometheus指标
request_counter = Counter(  # Counter统计元素出现次数
    'inference_requests_total',
    'Total inference requests',
    ['model', 'status']
)

request_duration = Histogram(
    'inference_request_duration_seconds',
    'Inference request duration',
    ['model']
)

gpu_memory_usage = Gauge(
    'gpu_memory_usage_bytes',
    'GPU memory usage'
)

cpu_usage = Gauge(
    'cpu_usage_percent',
    'CPU usage percentage'
)

active_conversations = Gauge(
    'active_conversations',
    'Number of active conversations'
)

@app.get("/metrics")
async def metrics():
    """
    Prometheus指标
    """
    # 更新系统指标
    if torch.cuda.is_available():
        gpu_memory_usage.set(torch.cuda.memory_allocated())

    cpu_usage.set(psutil.cpu_percent())

    return Response(
        content=generate_latest(),
        media_type=CONTENT_TYPE_LATEST
    )

@app.get("/health")
async def health_check():
    """
    健康检查
    """
    return {
        "status": "healthy",
        "gpu_available": torch.cuda.is_available(),
        "cpu_usage": psutil.cpu_percent()
    }

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