LLM应用开发——本地模型部署

import ollama
from typing import Generator

class OllamaClient:
    """Ollama 客户端封装"""
    
    def __init__(self, host: str = "http://localhost:11434"):
        self.client = ollama.Client(host=host)
    
    def chat(
        self,
        model: str,
        messages: list,
        stream: bool = False
    ) -> str | Generator:
        """对话"""
        response = self.client.chat(
            model=model,
            messages=messages,
            stream=stream
        )
        
        if stream:
            return self._stream_response(response)
        return response["message"]["content"]
    
    def _stream_response(self, response) -> Generator:
        """流式响应"""
        for chunk in response:
            yield chunk["message"]["content"]
    
    def generate(
        self,
        model: str,
        prompt: str,
        system: str = None
    ) -> str:
        """生成"""
        response = self.client.generate(
            model=model,
            prompt=prompt,
            system=system
        )
        return response["response"]
    
    def embed(self, model: str, text: str) -> list:
        """获取嵌入向量"""
        response = self.client.embeddings(
            model=model,
            prompt=text
        )
        return response["embedding"]
    
    def list_models(self) -> list:
        """列出模型"""
        response = self.client.list()
        return [m["name"] for m in response["models"]]

# 使用
client = OllamaClient()

# 对话
messages = [
    {"role": "user", "content": "用Python写一个快速排序"}
]
response = client.chat("llama3.1:8b", messages)
print(response)

# 流式对话
for chunk in client.chat("llama3.1:8b", messages, stream=True):
    print(chunk, end="", flush=True)

# 嵌入
embedding = client.embed("nomic-embed-text", "Hello world")

import requests
import json

class OllamaAPI:
    """Ollama REST API"""
    
    def __init__(self, base_url: str = "http://localhost:11434"):
        self.base_url = base_url
    
    def generate(self, model: str, prompt: str) -> str:
        """生成"""
        response = requests.post(
            f"{self.base_url}/api/generate",
            json={"model": model, "prompt": prompt, "stream": False}
        )
        return response.json()["response"]
    
    def chat_stream(self, model: str, messages: list):
        """流式对话"""
        response = requests.post(
            f"{self.base_url}/api/chat",
            json={"model": model, "messages": messages, "stream": True},
            stream=True
        )
        
        for line in response.iter_lines():
            if line:
                data = json.loads(line)
                if "message" in data:
                    yield data["message"]["content"]
    
    def pull_model(self, model: str):
        """拉取模型"""
        response = requests.post(
            f"{self.base_url}/api/pull",
            json={"name": model},
            stream=True
        )
        
        for line in response.iter_lines():
            if line:
                print(json.loads(line))

# 关键指标说明
# vllm:num_requests_running: 当前运行中的请求数
# vllm:iteration_tokens_total: 每秒生成的 Token 总数
# vllm:gpu_cache_usage_perc: KV Cache 显存占用率

from huggingface_hub import InferenceClient
from text_generation import Client

class TGIClient:
    """TGI 客户端"""
    
    def __init__(self, url: str = "http://localhost:8080"):
        self.client = Client(url)
        self.hf_client = InferenceClient(url)
    
    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 512,
        temperature: float = 0.7
    ) -> str:
        """生成"""
        response = self.client.generate(
            prompt,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=0.95,
            do_sample=True
        )
        return response.generated_text
    
    def stream(self, prompt: str, max_new_tokens: int = 512):
        """流式生成"""
        for token in self.client.generate_stream(
            prompt,
            max_new_tokens=max_new_tokens
        ):
            yield token.token.text
    
    def chat(self, messages: list) -> str:
        """对话"""
        response = self.hf_client.chat_completion(
            messages,
            max_tokens=512
        )
        return response.choices[0].message.content

# 使用
tgi_client = TGIClient()
response = tgi_client.generate("写一个Python函数计算斐波那契数列")

for token in tgi_client.stream("解释什么是机器学习"):
    print(token, end="", flush=True)

from llama_cpp import Llama

class LlamaCppClient:
    """llama.cpp Python 客户端"""
    
    def __init__(
        self,
        model_path: str,
        n_ctx: int = 4096,
        n_gpu_layers: int = -1  # -1 = 全部上 GPU
    ):
        self.llm = Llama(
            model_path=model_path,
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu_layers,
            chat_format="llama-2",
            verbose=False
        )
    
    def generate(
        self,
        prompt: str,
        max_tokens: int = 512,
        temperature: float = 0.7
    ) -> str:
        """生成"""
        output = self.llm(
            prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            stop=["</s>", "\n\n"]
        )
        return output["choices"][0]["text"]
    
    def chat(self, messages: list) -> str:
        """对话"""
        response = self.llm.create_chat_completion(
            messages,
            max_tokens=512
        )
        return response["choices"][0]["message"]["content"]
    
    def stream(self, prompt: str):
        """流式生成"""
        for chunk in self.llm(
            prompt,
            max_tokens=512,
            stream=True
        ):
            yield chunk["choices"][0]["text"]

# 使用
client = LlamaCppClient(
    "models/llama-7b-q4_k_m.gguf",
    n_gpu_layers=35
)

response = client.chat([
    {"role": "user", "content": "你好"}
])

from transformers import AutoModelForCausalLM, AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

def quantize_gptq(
    model_name: str,
    output_dir: str,
    bits: int = 4
):
    """GPTQ 量化"""
    # 加载模型
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # 量化配置
    quantize_config = BaseQuantizeConfig(
        bits=bits,
        group_size=128,
        desc_act=True
    )
    
    # 准备校准数据
    calibration_data = [
        "这是一段用于校准的文本...",
        "另一段校准文本..."
    ]
    
    # 加载并量化
    model = AutoGPTQForCausalLM.from_pretrained(
        model_name,
        quantize_config
    )
    
    model.quantize(calibration_data)
    model.save_quantized(output_dir)
    tokenizer.save_pretrained(output_dir)

# 使用量化模型
model = AutoGPTQForCausalLM.from_quantized(
    "model-gptq",
    device="cuda:0"
)

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

def quantize_awq(
    model_name: str,
    output_dir: str
):
    """AWQ 量化"""
    # 加载
    model = AutoAWQForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # 量化配置
    quant_config = {
        "zero_point": True,
        "q_group_size": 128,
        "w_bit": 4
    }
    
    # 量化
    model.quantize(tokenizer, quant_config=quant_config)
    
    # 保存
    model.save_quantized(output_dir)
    tokenizer.save_pretrained(output_dir)

# 使用
model = AutoAWQForCausalLM.from_quantized(
    "model-awq",
    fuse_layers=True
)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, List

app = FastAPI()

class ChatRequest(BaseModel):
    model: str
    messages: List[dict]
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 512
    stream: Optional[bool] = False

class LocalLLMRouter:
    """本地 LLM 路由器"""
    
    def __init__(self):
        self.backends = {}
    
    def register_backend(self, name: str, client):
        """注册后端"""
        self.backends[name] = client
    
    def chat(self, request: ChatRequest) -> str:
        """路由请求"""
        if request.model not in self.backends:
            raise ValueError(f"Model {request.model} not found")
        
        client = self.backends[request.model]
        return client.chat(request.messages)

router = LocalLLMRouter()

# 注册不同后端
router.register_backend("llama3", OllamaClient())
router.register_backend("codellama", VLLMServer("codellama/CodeLlama-7b"))

@app.post("/v1/chat/completions")
async def chat(request: ChatRequest):
    """OpenAI 兼容接口"""
    try:
        response = router.chat(request)
        return {
            "choices": [{
                "message": {"role": "assistant", "content": response}
            }]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))