Create openai_api_server_hf.py
无 vllm 环境下使用。 bug: 请求时，stream =true, 时会返回空，必须设置 stream =False 参考自： https://github.com/leisc/glm4_openai_api_server_mps
This commit is contained in:
mingyue0094
GitHub
parent
b562723d59
commit
fed14652f8
|
|
@ -0,0 +1,588 @@
|
|||
import os
|
||||
import time
|
||||
from asyncio.log import logger
|
||||
|
||||
import uvicorn
|
||||
import gc
|
||||
import json
|
||||
import torch
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Response
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import List, Literal, Optional, Union, Tuple
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessor
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
EventSourceResponse.DEFAULT_PING_INTERVAL = 10000
|
||||
MODEL_PATH = os.environ.get('MODEL_PATH', 'THU/glm-4-9b-chat')
|
||||
|
||||
TOKENIZER_PATH = os.environ.get("TOKENIZER_PATH", MODEL_PATH)
|
||||
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
yield
|
||||
if torch.cuda.is_available():
|
||||
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 FunctionCallResponse(BaseModel):
|
||||
name: Optional[str] = None
|
||||
arguments: Optional[str] = None
|
||||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
role: Literal["user", "assistant", "system", "tool"]
|
||||
content: str = None
|
||||
name: Optional[str] = None
|
||||
function_call: Optional[FunctionCallResponse] = None
|
||||
|
||||
|
||||
class DeltaMessage(BaseModel):
|
||||
role: Optional[Literal["user", "assistant", "system"]] = None
|
||||
content: Optional[str] = None
|
||||
function_call: Optional[FunctionCallResponse] = None
|
||||
|
||||
|
||||
class EmbeddingRequest(BaseModel):
|
||||
input: Union[List[str], str]
|
||||
model: str
|
||||
|
||||
|
||||
class CompletionUsage(BaseModel):
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
total_tokens: int
|
||||
|
||||
|
||||
class EmbeddingResponse(BaseModel):
|
||||
data: list
|
||||
model: str
|
||||
object: str
|
||||
usage: CompletionUsage
|
||||
|
||||
|
||||
class UsageInfo(BaseModel):
|
||||
prompt_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
completion_tokens: Optional[int] = 0
|
||||
|
||||
|
||||
class ChatCompletionRequest(BaseModel):
|
||||
model: str
|
||||
messages: List[ChatMessage]
|
||||
temperature: Optional[float] = 0.8
|
||||
top_p: Optional[float] = 0.8
|
||||
max_tokens: Optional[int] = None
|
||||
stream: Optional[bool] = False
|
||||
tools: Optional[Union[dict, List[dict]]] = None
|
||||
tool_choice: Optional[Union[str, dict]] = "None"
|
||||
repetition_penalty: Optional[float] = 1.1
|
||||
|
||||
|
||||
class ChatCompletionResponseChoice(BaseModel):
|
||||
index: int
|
||||
message: ChatMessage
|
||||
finish_reason: Literal["stop", "length", "function_call"]
|
||||
|
||||
|
||||
class ChatCompletionResponseStreamChoice(BaseModel):
|
||||
delta: DeltaMessage
|
||||
finish_reason: Optional[Literal["stop", "length", "function_call"]]
|
||||
index: int
|
||||
|
||||
|
||||
class ChatCompletionResponse(BaseModel):
|
||||
model: str
|
||||
id: str
|
||||
object: Literal["chat.completion", "chat.completion.chunk"]
|
||||
choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]]
|
||||
created: Optional[int] = Field(default_factory=lambda: int(time.time()))
|
||||
usage: Optional[UsageInfo] = None
|
||||
|
||||
|
||||
class InvalidScoreLogitsProcessor(LogitsProcessor):
|
||||
def __call__(
|
||||
self, input_ids: torch.LongTensor, scores: torch.FloatTensor
|
||||
) -> torch.FloatTensor:
|
||||
if torch.isnan(scores).any() or torch.isinf(scores).any():
|
||||
scores.zero_()
|
||||
scores[..., 5] = 5e4
|
||||
return scores
|
||||
|
||||
|
||||
def process_response(output: str, use_tool: bool = False) -> Union[str, dict]:
|
||||
content = ""
|
||||
for response in output.split(""):
|
||||
metadata, content = response.split("\n", maxsplit=1)
|
||||
if not metadata.strip():
|
||||
content = content.strip()
|
||||
else:
|
||||
if use_tool:
|
||||
content = "\n".join(content.split("\n")[1:-1])
|
||||
parameters = eval(content)
|
||||
content = {
|
||||
"name": metadata.strip(),
|
||||
"arguments": json.dumps(parameters, ensure_ascii=False)
|
||||
}
|
||||
else:
|
||||
content = {
|
||||
"name": metadata.strip(),
|
||||
"content": content
|
||||
}
|
||||
return content
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate_stream_glm4(params: dict):
|
||||
global engine, tokenizer
|
||||
|
||||
echo = params.get("echo", True)
|
||||
messages = params["messages"]
|
||||
tools = params["tools"]
|
||||
tool_choice = params["tool_choice"]
|
||||
temperature = float(params.get("temperature", 1.0))
|
||||
repetition_penalty = float(params.get("repetition_penalty", 1.0))
|
||||
top_p = float(params.get("top_p", 1.0))
|
||||
max_new_tokens = int(params.get("max_tokens", 8192))
|
||||
messages = process_messages(messages, tools=tools, tool_choice=tool_choice)
|
||||
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
|
||||
return_tensors="pt",return_dict=True)
|
||||
|
||||
inputs = inputs.to(engine.device)
|
||||
input_echo_len = len(inputs["input_ids"][0])
|
||||
|
||||
if input_echo_len >= engine.config.seq_length:
|
||||
print(f"Input length larger than {model.config.seq_length}")
|
||||
|
||||
eos_token_id = [tokenizer.eos_token_id,
|
||||
tokenizer.convert_tokens_to_ids("<|user|>"),
|
||||
tokenizer.convert_tokens_to_ids("<|observation|>")]
|
||||
|
||||
gen_kwargs = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"do_sample": True if temperature > 1e-5 else False,
|
||||
"top_p": top_p,
|
||||
"repetition_penalty": repetition_penalty,
|
||||
"logits_processor": [InvalidScoreLogitsProcessor()],
|
||||
}
|
||||
if temperature > 1e-5:
|
||||
gen_kwargs["temperature"] = temperature
|
||||
|
||||
total_len = 0
|
||||
for total_ids in engine.stream_generate(**inputs, eos_token_id=eos_token_id, **gen_kwargs):
|
||||
total_ids = total_ids.tolist()[0]
|
||||
total_len = len(total_ids)
|
||||
|
||||
if echo:
|
||||
output_ids = total_ids[:-1]
|
||||
else:
|
||||
output_ids = total_ids[input_echo_len:-1]
|
||||
|
||||
response = tokenizer.decode(output_ids)
|
||||
if response and response[-1] != "<EFBFBD>":
|
||||
response, stop_found = apply_stopping_strings(response, ["<|observation|>"])
|
||||
|
||||
yield {
|
||||
"text": response,
|
||||
"usage": {
|
||||
"prompt_tokens": input_echo_len,
|
||||
"completion_tokens": total_len - input_echo_len,
|
||||
"total_tokens": total_len,
|
||||
},
|
||||
"finish_reason": "function_call" if stop_found else None,
|
||||
}
|
||||
|
||||
if stop_found:
|
||||
break
|
||||
|
||||
# Only last stream result contains finish_reason, we set finish_reason as stop
|
||||
ret = {
|
||||
"text": response,
|
||||
"usage": {
|
||||
"prompt_tokens": input_echo_len,
|
||||
"completion_tokens": total_len - input_echo_len,
|
||||
"total_tokens": total_len,
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
yield ret
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def apply_stopping_strings(reply, stop_strings) -> Tuple[str, bool]:
|
||||
stop_found = False
|
||||
for string in stop_strings:
|
||||
idx = reply.find(string)
|
||||
if idx != -1:
|
||||
reply = reply[:idx]
|
||||
stop_found = True
|
||||
break
|
||||
|
||||
if not stop_found:
|
||||
# If something like "\nYo" is generated just before "\nYou: is completed, trim it
|
||||
for string in stop_strings:
|
||||
for j in range(len(string) - 1, 0, -1):
|
||||
if reply[-j:] == string[:j]:
|
||||
reply = reply[:-j]
|
||||
break
|
||||
else:
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
return reply, stop_found
|
||||
|
||||
|
||||
def process_messages(messages, tools=None, tool_choice="none"):
|
||||
_messages = messages
|
||||
messages = []
|
||||
msg_has_sys = False
|
||||
|
||||
def filter_tools(tool_choice, tools):
|
||||
function_name = tool_choice.get('function', {}).get('name', None)
|
||||
if not function_name:
|
||||
return []
|
||||
filtered_tools = [
|
||||
tool for tool in tools
|
||||
if tool.get('function', {}).get('name') == function_name
|
||||
]
|
||||
return filtered_tools
|
||||
|
||||
if tool_choice != "none":
|
||||
if isinstance(tool_choice, dict):
|
||||
tools = filter_tools(tool_choice, tools)
|
||||
if tools:
|
||||
messages.append(
|
||||
{
|
||||
"role": "system",
|
||||
"content": None,
|
||||
"tools": tools
|
||||
}
|
||||
)
|
||||
msg_has_sys = True
|
||||
|
||||
# add to metadata
|
||||
if isinstance(tool_choice, dict) and tools:
|
||||
messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"metadata": tool_choice["function"]["name"],
|
||||
"content": ""
|
||||
}
|
||||
)
|
||||
|
||||
for m in _messages:
|
||||
role, content, func_call = m.role, m.content, m.function_call
|
||||
if role == "function":
|
||||
messages.append(
|
||||
{
|
||||
"role": "observation",
|
||||
"content": content
|
||||
}
|
||||
)
|
||||
elif role == "assistant" and func_call is not None:
|
||||
for response in content.split(""):
|
||||
metadata, sub_content = response.split("\n", maxsplit=1)
|
||||
messages.append(
|
||||
{
|
||||
"role": role,
|
||||
"metadata": metadata,
|
||||
"content": sub_content.strip()
|
||||
}
|
||||
)
|
||||
else:
|
||||
if role == "system" and msg_has_sys:
|
||||
msg_has_sys = False
|
||||
continue
|
||||
messages.append({"role": role, "content": content})
|
||||
|
||||
return messages
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> Response:
|
||||
"""Health check."""
|
||||
return Response(status_code=200)
|
||||
|
||||
|
||||
@app.get("/v1/models", response_model=ModelList)
|
||||
async def list_models():
|
||||
model_card = ModelCard(id="glm-4")
|
||||
return ModelList(data=[model_card])
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
|
||||
async def create_chat_completion(request: ChatCompletionRequest):
|
||||
if len(request.messages) < 1 or request.messages[-1].role == "assistant":
|
||||
raise HTTPException(status_code=400, detail="Invalid request")
|
||||
|
||||
gen_params = dict(
|
||||
messages=request.messages,
|
||||
temperature=request.temperature,
|
||||
top_p=request.top_p,
|
||||
max_tokens=request.max_tokens or 1024,
|
||||
echo=False,
|
||||
stream=request.stream,
|
||||
repetition_penalty=request.repetition_penalty,
|
||||
tools=request.tools,
|
||||
tool_choice=request.tool_choice,
|
||||
)
|
||||
logger.debug(f"==== request ====\n{gen_params}")
|
||||
|
||||
if request.stream:
|
||||
predict_stream_generator = predict_stream(request.model, gen_params)
|
||||
output = await anext(predict_stream_generator)
|
||||
if not output and 'get_' in output:
|
||||
return EventSourceResponse(predict_stream_generator, media_type="text/event-stream")
|
||||
logger.debug(f"First result output：\n{output}")
|
||||
|
||||
function_call = None
|
||||
if output and request.tools:
|
||||
try:
|
||||
function_call = process_response(output, use_tool=True)
|
||||
except:
|
||||
logger.warning("Failed to parse tool call")
|
||||
|
||||
# CallFunction
|
||||
if isinstance(function_call, dict):
|
||||
function_call = FunctionCallResponse(**function_call)
|
||||
tool_response = ""
|
||||
if not gen_params.get("messages"):
|
||||
gen_params["messages"] = []
|
||||
gen_params["messages"].append(ChatMessage(role="assistant", content=output))
|
||||
gen_params["messages"].append(ChatMessage(role="tool", name=function_call.name, content=tool_response))
|
||||
generate = predict(request.model, gen_params)
|
||||
return EventSourceResponse(generate, media_type="text/event-stream")
|
||||
else:
|
||||
generate = parse_output_text(request.model, output)
|
||||
return EventSourceResponse(generate, media_type="text/event-stream")
|
||||
|
||||
response = ""
|
||||
for response in generate_stream_glm4(gen_params):
|
||||
pass
|
||||
|
||||
if response["text"].startswith("\n"):
|
||||
response["text"] = response["text"][1:]
|
||||
response["text"] = response["text"].strip()
|
||||
|
||||
usage = UsageInfo()
|
||||
function_call, finish_reason = None, "stop"
|
||||
if request.tools:
|
||||
try:
|
||||
function_call = process_response(response["text"], use_tool=True)
|
||||
except:
|
||||
logger.warning(
|
||||
"Failed to parse tool call, maybe the response is not a function call(such as cogview drawing) or have been answered.")
|
||||
|
||||
if isinstance(function_call, dict):
|
||||
finish_reason = "function_call"
|
||||
function_call = FunctionCallResponse(**function_call)
|
||||
|
||||
message = ChatMessage(
|
||||
role="assistant",
|
||||
content=response["text"],
|
||||
function_call=function_call if isinstance(function_call, FunctionCallResponse) else None,
|
||||
)
|
||||
|
||||
logger.debug(f"==== message ====\n{message}")
|
||||
|
||||
choice_data = ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=message,
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
task_usage = UsageInfo.model_validate(response["usage"])
|
||||
for usage_key, usage_value in task_usage.model_dump().items():
|
||||
setattr(usage, usage_key, getattr(usage, usage_key) + usage_value)
|
||||
|
||||
return ChatCompletionResponse(
|
||||
model=request.model,
|
||||
id="", # for open_source model, id is empty
|
||||
choices=[choice_data],
|
||||
object="chat.completion",
|
||||
usage=usage
|
||||
)
|
||||
|
||||
|
||||
async def predict(model_id: str, params: dict):
|
||||
choice_data = ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role="assistant"),
|
||||
finish_reason=None
|
||||
)
|
||||
chunk = ChatCompletionResponse(model=model_id, id="", choices=[choice_data], object="chat.completion.chunk")
|
||||
yield "{}".format(chunk.model_dump_json(exclude_unset=True))
|
||||
|
||||
previous_text = ""
|
||||
for new_response in generate_stream_glm4(params):
|
||||
decoded_unicode = new_response["text"]
|
||||
delta_text = decoded_unicode[len(previous_text):]
|
||||
previous_text = decoded_unicode
|
||||
|
||||
finish_reason = new_response["finish_reason"]
|
||||
if len(delta_text) == 0 and finish_reason != "function_call":
|
||||
continue
|
||||
|
||||
function_call = None
|
||||
if finish_reason == "function_call":
|
||||
try:
|
||||
function_call = process_response(decoded_unicode, use_tool=True)
|
||||
except:
|
||||
logger.warning(
|
||||
"Failed to parse tool call, maybe the response is not a tool call or have been answered.")
|
||||
|
||||
if isinstance(function_call, dict):
|
||||
function_call = FunctionCallResponse(**function_call)
|
||||
|
||||
delta = DeltaMessage(
|
||||
content=delta_text,
|
||||
role="assistant",
|
||||
function_call=function_call if isinstance(function_call, FunctionCallResponse) else None,
|
||||
)
|
||||
|
||||
choice_data = ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=delta,
|
||||
finish_reason=finish_reason
|
||||
)
|
||||
chunk = ChatCompletionResponse(
|
||||
model=model_id,
|
||||
id="",
|
||||
choices=[choice_data],
|
||||
object="chat.completion.chunk"
|
||||
)
|
||||
yield "{}".format(chunk.model_dump_json(exclude_unset=True))
|
||||
|
||||
choice_data = ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(),
|
||||
finish_reason="stop"
|
||||
)
|
||||
chunk = ChatCompletionResponse(
|
||||
model=model_id,
|
||||
id="",
|
||||
choices=[choice_data],
|
||||
object="chat.completion.chunk"
|
||||
)
|
||||
yield "{}".format(chunk.model_dump_json(exclude_unset=True))
|
||||
yield '[DONE]'
|
||||
|
||||
|
||||
async def predict_stream(model_id, gen_params):
|
||||
output = ""
|
||||
is_function_call = False
|
||||
has_send_first_chunk = False
|
||||
for new_response in generate_stream_glm4(gen_params):
|
||||
decoded_unicode = new_response["text"]
|
||||
delta_text = decoded_unicode[len(output):]
|
||||
output = decoded_unicode
|
||||
|
||||
if not is_function_call and len(output) > 7:
|
||||
is_function_call = output and 'get_' in output
|
||||
if is_function_call:
|
||||
continue
|
||||
|
||||
finish_reason = new_response["finish_reason"]
|
||||
if not has_send_first_chunk:
|
||||
message = DeltaMessage(
|
||||
content="",
|
||||
role="assistant",
|
||||
function_call=None,
|
||||
)
|
||||
choice_data = ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=message,
|
||||
finish_reason=finish_reason
|
||||
)
|
||||
chunk = ChatCompletionResponse(
|
||||
model=model_id,
|
||||
id="",
|
||||
choices=[choice_data],
|
||||
created=int(time.time()),
|
||||
object="chat.completion.chunk"
|
||||
)
|
||||
yield "{}".format(chunk.model_dump_json(exclude_unset=True))
|
||||
|
||||
send_msg = delta_text if has_send_first_chunk else output
|
||||
has_send_first_chunk = True
|
||||
message = DeltaMessage(
|
||||
content=send_msg,
|
||||
role="assistant",
|
||||
function_call=None,
|
||||
)
|
||||
choice_data = ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=message,
|
||||
finish_reason=finish_reason
|
||||
)
|
||||
chunk = ChatCompletionResponse(
|
||||
model=model_id,
|
||||
id="",
|
||||
choices=[choice_data],
|
||||
created=int(time.time()),
|
||||
object="chat.completion.chunk"
|
||||
)
|
||||
yield "{}".format(chunk.model_dump_json(exclude_unset=True))
|
||||
|
||||
if is_function_call:
|
||||
yield output
|
||||
else:
|
||||
yield '[DONE]'
|
||||
|
||||
|
||||
async def parse_output_text(model_id: str, value: str):
|
||||
choice_data = ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role="assistant", content=value),
|
||||
finish_reason=None
|
||||
)
|
||||
chunk = ChatCompletionResponse(model=model_id, id="", choices=[choice_data], object="chat.completion.chunk")
|
||||
yield "{}".format(chunk.model_dump_json(exclude_unset=True))
|
||||
choice_data = ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(),
|
||||
finish_reason="stop"
|
||||
)
|
||||
chunk = ChatCompletionResponse(model=model_id, id="", choices=[choice_data], object="chat.completion.chunk")
|
||||
yield "{}".format(chunk.model_dump_json(exclude_unset=True))
|
||||
yield '[DONE]'
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True)
|
||||
|
||||
engine = AutoModelForCausalLM.from_pretrained(MODEL_PATH,trust_remote_code=True,load_in_4bit=True,torch_dtype=torch.bfloat16,device_map="cuda",).eval() # 4bit 量化
|
||||
#engine = AutoModelForCausalLM.from_pretrained(MODEL_PATH,trust_remote_code=True,torch_dtype=torch.bfloat16,device_map="cuda",).eval() #无量化
|
||||
uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) #请求时,必须 stream=False， 否则返回空内容。
|
||||
Loading…
Reference in New Issue