# coding=utf-8 # Implements API for ChatGLM3-6B in OpenAI's format. (https://platform.openai.com/docs/api-reference/chat) # Usage: python openai_api.py # Visit http://localhost:8100/docs for documents. import time import json import torch import uvicorn import argparse from pydantic import BaseModel, Field from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from contextlib import asynccontextmanager from typing import Any, Dict, List, Literal, Optional, Union #from transformers import AutoTokenizer, AutoModel from sse_starlette.sse import ServerSentEvent, EventSourceResponse from fastllm_pytools import llm @asynccontextmanager async def lifespan(app: FastAPI): # collects GPU memory yield global device_map if torch.cuda.is_available(): for device in device_map: with torch.cuda.device(device): torch.cuda.empty_cache() torch.cuda.ipc_collect() app = FastAPI(lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) class ModelCard(BaseModel): id: str object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "owner" root: Optional[str] = None parent: Optional[str] = None permission: Optional[list] = None class ModelList(BaseModel): object: str = "list" data: List[ModelCard] = [] class ChatMessage(BaseModel): role: Literal["user", "assistant", "system"] content: str class Usage(BaseModel): prompt_tokens: int = None total_tokens: int = None completion_tokens: int = None class DeltaMessage(BaseModel): role: Optional[Literal["user", "assistant", "system"]] = None content: Optional[str] = None class ChatCompletionRequest(BaseModel): model: str messages: List[ChatMessage] temperature: Optional[float] = None top_p: Optional[float] = None max_length: Optional[int] = None stream: Optional[bool] = False class ChatCompletionResponseChoice(BaseModel): index: int message: ChatMessage finish_reason: Literal["stop", "length"] class ChatCompletionResponseStreamChoice(BaseModel): index: int delta: DeltaMessage finish_reason: Optional[Literal["stop", "length"]] class ChatCompletionResponse(BaseModel): id: str object: Literal["chat.completion", "chat.completion.chunk"] created: Optional[int] = Field(default_factory=lambda: int(time.time())) model: str choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]] usage: Usage = None @app.get("/v1/models", response_model=ModelList) def list_models(): global model_list for model in model_list: ModelCard(id=model) ModelList.data.append(ModelCard) return ModelList() @app.post("/v1/chat/completions", response_model=ChatCompletionResponse) def create_chat_completion(request: ChatCompletionRequest): if request.model not in model_list: raise HTTPException(status_code=400, detail="Invalid Model Name") global model id = "chatcmpl-A" if request.messages[-1].role != "user": raise HTTPException(status_code=400, detail="Invalid request") query = request.messages[-1].content if request.max_length is not None: max_length = request.max_length else: max_length = 1024 if request.temperature is not None: temperature = request.temperature else: temperature = 0.1 if request.top_p is not None: top_p = request.top_p else: top_p = 0.8 prev_messages = request.messages[:-1] # print(prev_messages) if len(prev_messages) > 0 and prev_messages[0].role == "system": query = prev_messages.pop(0).content + query history = [] if len(prev_messages) % 2 == 0: for i in range(0, len(prev_messages), 2): if prev_messages[i].role == "user" and prev_messages[i+1].role == "assistant": history.append([prev_messages[i].content, prev_messages[i+1].content]) if request.stream: generate = predict(id=id, query=query, history=history, max_length=max_length, top_p = top_p, temperature = temperature, model_id = request.model) return EventSourceResponse(generate, media_type="text/event-stream") response = model.response(query=query, history=history, max_length=max_length, top_p = top_p, temperature = temperature) choice_data = ChatCompletionResponseChoice( index=0, message=ChatMessage(role="assistant", content=response), finish_reason="stop" ) prompt_tokens = len(model.tokenizer_encode_string(query)) completion_tokens = len(model.tokenizer_encode_string(response)) usage = Usage( prompt_tokens = prompt_tokens, completion_tokens = completion_tokens, total_tokens = prompt_tokens+completion_tokens, ) return ChatCompletionResponse(id=id ,model=request.model, choices=[choice_data], object="chat.completion", usage=usage) def predict(id: str, query: str, history: List[List[str]], model_id: str, max_length: int, top_p: float, temperature: float): global model creat_time = int(time.time()) choice_data = ChatCompletionResponseStreamChoice( index=0, delta=DeltaMessage(role="assistant"), finish_reason=None ) chunk = ChatCompletionResponse(id=id, created=creat_time, model=model_id, choices=[choice_data], object="chat.completion.chunk") #yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False)) //pydantic从1.8.0开始不支持dumps_kwags参数，参考https://github.com/THUDM/ChatGLM2-6B/issues/308 yield json.dumps(chunk.model_dump(exclude_unset=True), ensure_ascii=False) for new_response in model.stream_response(query=query, history=history, max_length=max_length, top_p = top_p, temperature = temperature): choice_data = ChatCompletionResponseStreamChoice( index=0, delta=DeltaMessage(content=new_response), finish_reason=None ) chunk = ChatCompletionResponse(id=id, created=creat_time, model=model_id, choices=[choice_data], object="chat.completion.chunk") #yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False)) yield json.dumps(chunk.model_dump(exclude_unset=True), ensure_ascii=False) choice_data = ChatCompletionResponseStreamChoice( index=0, delta=DeltaMessage(), finish_reason="stop" ) chunk = ChatCompletionResponse(id=id, created=creat_time, model=model_id, choices=[choice_data], object="chat.completion.chunk") #yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False)) yield json.dumps(chunk.model_dump(exclude_unset=True), ensure_ascii=False) yield '[DONE]' def args_parser(): parser = argparse.ArgumentParser(description = 'baichuan2_chat_demo') parser.add_argument('-p', '--path', type = str, default = "/model", help = '模型文件的路径') parser.add_argument('-g', '--gpus', type = str, default = "0", help = '指定运行的gpu卡，例如“0，1”') args = parser.parse_args() return args if __name__ == "__main__": args = args_parser() global model_list model_list = ["chatglm3-6b-fastllm"] global device_map device_map = ["cuda:"+num for num in args.gpus.split(',')] llm.set_device_map(device_map) model = llm.model(args.path) uvicorn.run(app, host='127.0.0.1', port=8100)