Fine-tuning/backend/app/core/deploy_server_template.py

"""OpenAI 兼容的模型推理服务器。

使用方法:
    python server.py --port 8000 --host 0.0.0.0

API 端点:
    POST /v1/chat/completions   - OpenAI 兼容的聊天补全接口
    POST /v1/completions        - 文本补全接口
    GET  /v1/models             - 模型列表
    GET  /health                - 健康检查

调用示例:
    curl http://localhost:8000/v1/chat/completions \
        -H "Content-Type: application/json" \
        -d '{
            "model": "local-model",
            "messages": [{"role": "user", "content": "你好"}],
            "max_tokens": 512,
            "temperature": 0.7
        }'
"""
import argparse
import json
import time
import uuid
from pathlib import Path

import torch
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList

app = FastAPI(title="Model Serving API", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

model = None
tokenizer = None
model_name = "local-model"


# --- Request / Response schemas ---

class Message(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    model: str = "local-model"
    messages: list[Message]
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.9
    stream: bool = False


class CompletionRequest(BaseModel):
    model: str = "local-model"
    prompt: str
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.9
    stream: bool = False


class ChoiceMessage(BaseModel):
    role: str = "assistant"
    content: str


class Choice(BaseModel):
    index: int = 0
    message: ChoiceMessage
    finish_reason: str = "stop"


class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class ChatResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: list[Choice]
    usage: Usage


class CompletionChoice(BaseModel):
    index: int = 0
    text: str
    finish_reason: str = "stop"


class CompletionResponse(BaseModel):
    id: str
    object: str = "text_completion"
    created: int
    model: str
    choices: list[CompletionChoice]
    usage: Usage


class ModelInfo(BaseModel):
    id: str
    object: str = "model"
    created: int = 0
    owned_by: str = "local"


class ModelList(BaseModel):
    object: str = "list"
    data: list[ModelInfo]


# --- Endpoints ---

@app.post("/v1/chat/completions", response_model=ChatResponse)
async def chat_completions(req: ChatRequest):
    prompt = _build_prompt(req.messages)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    prompt_tokens = inputs["input_ids"].shape[1]

    stopping_criteria = _build_stop_criteria()

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=req.max_tokens,
            temperature=max(req.temperature, 0.01),
            top_p=req.top_p,
            do_sample=req.temperature > 0,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            stopping_criteria=stopping_criteria,
        )

    generated = tokenizer.decode(outputs[0][prompt_tokens:], skip_special_tokens=True)
    generated = _clean_generated(generated)
    completion_tokens = outputs.shape[1] - prompt_tokens

    return ChatResponse(
        id=f"chatcmpl-{uuid.uuid4().hex[:12]}",
        created=int(time.time()),
        model=req.model,
        choices=[Choice(message=ChoiceMessage(content=generated))],
        usage=Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        ),
    )


@app.post("/v1/completions", response_model=CompletionResponse)
async def completions(req: CompletionRequest):
    inputs = tokenizer(req.prompt, return_tensors="pt").to(model.device)
    prompt_tokens = inputs["input_ids"].shape[1]

    stopping_criteria = _build_stop_criteria()

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=req.max_tokens,
            temperature=max(req.temperature, 0.01),
            top_p=req.top_p,
            do_sample=req.temperature > 0,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            stopping_criteria=stopping_criteria,
        )

    generated = tokenizer.decode(outputs[0][prompt_tokens:], skip_special_tokens=True)
    generated = _clean_generated(generated)
    completion_tokens = outputs.shape[1] - prompt_tokens

    return CompletionResponse(
        id=f"cmpl-{uuid.uuid4().hex[:12]}",
        created=int(time.time()),
        model=req.model,
        choices=[CompletionChoice(text=generated)],
        usage=Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        ),
    )


@app.get("/v1/models", response_model=ModelList)
async def list_models():
    return ModelList(data=[ModelInfo(id=model_name)])


@app.get("/health")
async def health():
    return {"status": "ok", "model": model_name}


def _build_prompt(messages: list[Message]) -> str:
    """将 OpenAI 消息格式转为模型输入文本。"""
    parts = []
    for msg in messages:
        if msg.role == "system":
            parts.append(f"<|system|>\n{msg.content}")
        elif msg.role == "user":
            parts.append(f"<|user|>\n{msg.content}")
        elif msg.role == "assistant":
            parts.append(f"<|assistant|>\n{msg.content}")
    parts.append("<|assistant|>\n")
    return "\n".join(parts)


def _build_stop_criteria():
    """构建 StoppingCriteria，遇到角色切换标记时停止生成，防止复读。"""
    stop_phrases = ["<|user|>", "<|system|>", "<|assistant|>"]
    stop_token_ids = []
    for phrase in stop_phrases:
        ids = tokenizer.encode(phrase, add_special_tokens=False)
        stop_token_ids.append(ids)

    class StopOnRoleToken(StoppingCriteria):
        def __init__(self, stop_sequences):
            self.stop_sequences = stop_sequences

        def __call__(self, input_ids, scores, **kwargs):
            gen_seq = input_ids[0].tolist()
            for stop_ids in self.stop_sequences:
                if len(gen_seq) >= len(stop_ids):
                    if gen_seq[-len(stop_ids):] == stop_ids:
                        return True
            return False

    return StoppingCriteriaList([StopOnRoleToken(stop_token_ids)])


def _clean_generated(generated: str) -> str:
    """清理可能残留的角色标记。"""
    for marker in ["<|user|>", "<|system|>", "<|assistant|>"]:
        if marker in generated:
            generated = generated[:generated.index(marker)]
    return generated.strip()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Model Serving API Server")
    parser.add_argument("--model-path", type=str, default="./model", help="模型目录路径")
    parser.add_argument("--host", type=str, default="0.0.0.0", help="监听地址")
    parser.add_argument("--port", type=int, default=8000, help="监听端口")
    parser.add_argument("--device", type=str, default="auto", help="设备 (auto/cuda/cpu)")
    args = parser.parse_args()

    model_path = Path(args.model_path)
    model_name = model_path.name
    print(f"Loading model from: {model_path}")

    if args.device == "auto":
        # device_map="auto" 自动将模型层分散到所有可见 GPU
        device_map = "auto" if torch.cuda.is_available() else "cpu"
    elif args.device == "cuda":
        device_map = "auto"  # 自动多卡分散
    else:
        device_map = "cpu"

    tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        str(model_path), torch_dtype=torch.float16, device_map=device_map,
    )
    model.eval()
    print(f"Model loaded. Starting server on {args.host}:{args.port}")

    import uvicorn
    uvicorn.run(app, host=args.host, port=args.port)