import traceback
from abc import ABC
from typing import (
Optional,
List,
Union,
Tuple,
Dict,
Iterator,
Any,
)
import torch
from fastapi.responses import JSONResponse
from loguru import logger
from openai.types.chat import (
ChatCompletionMessage,
ChatCompletion,
ChatCompletionChunk,
)
from openai.types.chat import ChatCompletionMessageParam
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_chunk import (
ChoiceDelta,
ChoiceDeltaFunctionCall,
ChoiceDeltaToolCall,
)
from openai.types.chat.chat_completion_message import FunctionCall
from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall
from openai.types.completion import Completion
from openai.types.completion_choice import CompletionChoice, Logprobs
from openai.types.completion_usage import CompletionUsage
from transformers import PreTrainedModel, PreTrainedTokenizer
from api.adapter import get_prompt_adapter
from api.generation import (
build_baichuan_chat_input,
check_is_baichuan,
generate_stream_chatglm,
check_is_chatglm,
generate_stream_chatglm_v3,
build_qwen_chat_input,
check_is_qwen,
generate_stream,
build_xverse_chat_input,
check_is_xverse,
)
from api.generation.utils import get_context_length
from api.utils.compat import model_parse
from api.utils.constants import ErrorCode
from api.utils.request import create_error_response
server_error_msg = (
"**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
)
class DefaultEngine(ABC):
""" 基于原生 transformers 实现的模型引擎 """
def __init__(
self,
model: PreTrainedModel,
tokenizer: PreTrainedTokenizer,
device: Union[str, torch.device],
model_name: str,
context_len: Optional[int] = None,
prompt_name: Optional[str] = None,
use_streamer_v2: Optional[bool] = False,
):
"""
Initialize the Default class.
Args:
model (PreTrainedModel): The pre-trained model.
tokenizer (PreTrainedTokenizer): The tokenizer for the model.
device (Union[str, torch.device]): The device to use for inference.
model_name (str): The name of the model.
context_len (Optional[int], optional): The length of the context. Defaults to None.
prompt_name (Optional[str], optional): The name of the prompt. Defaults to None.
use_streamer_v2 (Optional[bool], optional): Whether to use Streamer V2. Defaults to False.
"""
self.model = model
self.tokenizer = tokenizer
self.device = model.device if hasattr(model, "device") else device
self.model_name = model_name.lower()
self.prompt_name = prompt_name.lower() if prompt_name is not None else None
self.context_len = context_len
self.use_streamer_v2 = use_streamer_v2
self.prompt_adapter = get_prompt_adapter(self.model_name, prompt_name=self.prompt_name)
self._prepare_for_generate()
self._fix_tokenizer()
def _prepare_for_generate(self):
"""
Prepare the object for text generation.
1. Sets the appropriate generate stream function based on the model name and type.
2. Updates the context length if necessary.
3. Checks and constructs the prompt.
4. Sets the context length if it is not already set.
"""
self.generate_stream_func = generate_stream
if "chatglm3" in self.model_name:
self.generate_stream_func = generate_stream_chatglm_v3
self.use_streamer_v2 = False
elif check_is_chatglm(self.model):
self.generate_stream_func = generate_stream_chatglm
elif check_is_qwen(self.model):
self.context_len = 8192 if self.context_len is None else self.context_len
self._check_construct_prompt()
if self.context_len is None:
self.context_len = get_context_length(self.model.config)
def _check_construct_prompt(self):
""" Check whether to need to construct prompts or inputs. """
self.construct_prompt = self.prompt_name is not None
if "chatglm3" in self.model_name:
logger.info("Using ChatGLM3 Model for Chat!")
elif check_is_baichuan(self.model):
logger.info("Using Baichuan Model for Chat!")
elif check_is_qwen(self.model):
logger.info("Using Qwen Model for Chat!")
elif check_is_xverse(self.model):
logger.info("Using Xverse Model for Chat!")
else:
self.construct_prompt = True
def _fix_tokenizer(self):
"""
Fix the tokenizer by adding the end-of-sequence (eos) token
and the padding (pad) token if they are missing.
"""
if self.tokenizer.eos_token_id is None:
self.tokenizer.eos_token = "<|endoftext|>"
logger.info(f"Add eos token: {self.tokenizer.eos_token}")
if self.tokenizer.pad_token_id is None:
if self.tokenizer.unk_token_id is not None:
self.tokenizer.pad_token = self.tokenizer.unk_token
else:
self.tokenizer.pad_token = self.tokenizer.eos_token
logger.info(f"Add pad token: {self.tokenizer.pad_token}")
def convert_to_inputs(
self,
prompt_or_messages: Union[List[ChatCompletionMessageParam], str],
infilling: Optional[bool] = False,
suffix_first: Optional[bool] = False,
**kwargs,
) -> Tuple[Union[List[int], Dict[str, Any]], Union[List[ChatCompletionMessageParam], str]]:
"""
Convert the prompt or messages into input format for the model.
Args:
prompt_or_messages: The prompt or messages to be converted.
infilling: Whether to perform infilling.
suffix_first: Whether to append the suffix first.
**kwargs: Additional keyword arguments.
Returns:
Tuple containing the converted inputs and the prompt or messages.
"""
# for completion
if isinstance(prompt_or_messages, str):
if infilling:
inputs = self.tokenizer(
prompt_or_messages, suffix_first=suffix_first,
).input_ids
elif check_is_qwen(self.model):
inputs = self.tokenizer(
prompt_or_messages, allowed_special="all", disallowed_special=()
).input_ids
elif check_is_chatglm(self.model):
inputs = self.tokenizer([prompt_or_messages], return_tensors="pt")
else:
inputs = self.tokenizer(prompt_or_messages).input_ids
if isinstance(inputs, list):
max_src_len = self.context_len - kwargs.get("max_tokens", 256) - 1
inputs = inputs[-max_src_len:]
else:
inputs, prompt_or_messages = self.apply_chat_template(prompt_or_messages, **kwargs)
return inputs, prompt_or_messages
def apply_chat_template(
self,
messages: List[ChatCompletionMessageParam],
max_new_tokens: Optional[int] = 256,
functions: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
tools: Optional[List[Dict[str, Any]]] = None,
**kwargs,
) -> Tuple[Union[List[int], Dict[str, Any]], Optional[str]]:
"""
Apply chat template to generate model inputs and prompt.
Args:
messages (List[ChatCompletionMessageParam]): List of chat completion message parameters.
max_new_tokens (Optional[int], optional): Maximum number of new tokens to generate. Defaults to 256.
functions (Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], optional): Functions to apply to the messages. Defaults to None.
tools (Optional[List[Dict[str, Any]]], optional): Tools to apply to the messages. Defaults to None.
**kwargs: Additional keyword arguments.
Returns:
Tuple[Union[List[int], Dict[str, Any]], Union[str, None]]: Tuple containing the generated inputs and prompt.
"""
if self.prompt_adapter.function_call_available:
messages = self.prompt_adapter.postprocess_messages(
messages, functions, tools=tools,
)
if functions or tools:
logger.debug(f"==== Messages with tools ====\n{messages}")
if self.construct_prompt:
prompt = self.prompt_adapter.apply_chat_template(messages)
if check_is_qwen(self.model):
inputs = self.tokenizer(prompt, allowed_special="all", disallowed_special=()).input_ids
elif check_is_chatglm(self.model):
inputs = self.tokenizer([prompt], return_tensors="pt")
else:
inputs = self.tokenizer(prompt).input_ids
if isinstance(inputs, list):
max_src_len = self.context_len - max_new_tokens - 1
inputs = inputs[-max_src_len:]
return inputs, prompt
else:
inputs = self.build_chat_inputs(
messages, max_new_tokens, functions, tools, **kwargs
)
return inputs, None
def build_chat_inputs(
self,
messages: List[ChatCompletionMessageParam],
max_new_tokens: Optional[int] = 256,
functions: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
tools: Optional[List[Dict[str, Any]]] = None,
**kwargs: Any,
) -> List[int]:
if "chatglm3" in self.model_name:
query, role = messages[-1]["content"], messages[-1]["role"]
inputs = self.tokenizer.build_chat_input(query, history=messages[:-1], role=role)
elif check_is_baichuan(self.model):
inputs = build_baichuan_chat_input(
self.tokenizer, messages, self.context_len, max_new_tokens
)
elif check_is_qwen(self.model):
inputs = build_qwen_chat_input(
self.tokenizer, messages, self.context_len, max_new_tokens, functions, tools,
)
elif check_is_xverse(self.model):
inputs = build_xverse_chat_input(
self.tokenizer, messages, self.context_len, max_new_tokens
)
else:
raise NotImplementedError
return inputs
def _generate(self, params: Dict[str, Any]) -> Iterator:
"""
Generates text based on the given parameters.
Args:
params (Dict[str, Any]): A dictionary containing the parameters for text generation.
Yields:
Iterator: A dictionary containing the generated text and error code.
"""
prompt_or_messages = params.get("prompt_or_messages")
inputs, prompt = self.convert_to_inputs(
prompt_or_messages,
infilling=params.get("infilling", False),
suffix_first=params.get("suffix_first", False),
max_new_tokens=params.get("max_tokens", 256),
functions=params.get("functions"),
tools=params.get("tools"),
)
params.update(dict(inputs=inputs, prompt=prompt))
try:
for output in self.generate_stream_func(self.model, self.tokenizer, params):
output["error_code"] = 0
yield output
except torch.cuda.OutOfMemoryError as e:
yield {
"text": f"{server_error_msg}\n\n({e})",
"error_code": ErrorCode.CUDA_OUT_OF_MEMORY,
}
except (ValueError, RuntimeError) as e:
traceback.print_exc()
yield {
"text": f"{server_error_msg}\n\n({e})",
"error_code": ErrorCode.INTERNAL_ERROR,
}
def _create_completion_stream(self, params: Dict[str, Any]) -> Iterator:
"""
Generates a stream of completions based on the given parameters.
Args:
params (Dict[str, Any]): The parameters for generating completions.
Yields:
Iterator: A stream of completion objects.
"""
for output in self._generate(params):
if output["error_code"] != 0:
yield output
return
logprobs = None
if params.get("logprobs") and output["logprobs"]:
logprobs = model_parse(Logprobs, output["logprobs"])
choice = CompletionChoice(
index=0,
text=output["delta"],
finish_reason="stop",
logprobs=logprobs,
)
yield Completion(
id=output["id"],
choices=[choice],
created=output["created"],
model=output["model"],
object="text_completion",
)
def _create_completion(self, params: Dict[str, Any]) -> Union[Completion, JSONResponse]:
"""
Creates a completion based on the given parameters.
Args:
params (Dict[str, Any]): The parameters for creating the completion.
Returns:
Completion: The generated completion object.
"""
last_output = None
for output in self._generate(params):
last_output = output
if last_output["error_code"] != 0:
return create_error_response(last_output["error_code"], last_output["text"])
logprobs = None
if params.get("logprobs") and last_output["logprobs"]:
logprobs = model_parse(Logprobs, last_output["logprobs"])
choice = CompletionChoice(
index=0,
text=last_output["text"],
finish_reason="stop",
logprobs=logprobs,
)
usage = model_parse(CompletionUsage, last_output["usage"])
return Completion(
id=last_output["id"],
choices=[choice],
created=last_output["created"],
model=last_output["model"],
object="text_completion",
usage=usage,
)
def _create_chat_completion_stream(self, params: Dict[str, Any]) -> Iterator:
"""
Creates a chat completion stream.
Args:
params (Dict[str, Any]): The parameters for generating the chat completion.
Yields:
Dict[str, Any]: The output of the chat completion stream.
"""
_id, _created, _model = None, None, None
has_function_call = False
for i, output in enumerate(self._generate(params)):
if output["error_code"] != 0:
yield output
return
_id, _created, _model = output["id"], output["created"], output["model"]
if i == 0:
choice = ChunkChoice(
index=0,
delta=ChoiceDelta(role="assistant", content=""),
finish_reason=None,
logprobs=None,
)
yield ChatCompletionChunk(
id=f"chat{_id}",
choices=[choice],
created=_created,
model=_model,
object="chat.completion.chunk",
)
finish_reason = output["finish_reason"]
if len(output["delta"]) == 0 and finish_reason != "function_call":
continue
function_call = None
if finish_reason == "function_call":
try:
_, function_call = self.prompt_adapter.parse_assistant_response(
output["text"], params.get("functions"), params.get("tools"),
)
except Exception as e:
traceback.print_exc()
logger.warning("Failed to parse tool call")
if isinstance(function_call, dict) and "arguments" in function_call:
has_function_call = True
function_call = ChoiceDeltaFunctionCall(**function_call)
delta = ChoiceDelta(
content=output["delta"],
function_call=function_call
)
elif isinstance(function_call, dict) and "function" in function_call:
has_function_call = True
finish_reason = "tool_calls"
function_call["index"] = 0
tool_calls = [model_parse(ChoiceDeltaToolCall, function_call)]
delta = ChoiceDelta(
content=output["delta"],
tool_calls=tool_calls,
)
else:
delta = ChoiceDelta(content=output["delta"])
choice = ChunkChoice(
index=0,
delta=delta,
finish_reason=finish_reason,
logprobs=None,
)
yield ChatCompletionChunk(
id=f"chat{_id}",
choices=[choice],
created=_created,
model=_model,
object="chat.completion.chunk",
)
if not has_function_call:
choice = ChunkChoice(
index=0,
delta=ChoiceDelta(),
finish_reason="stop",
logprobs=None,
)
yield ChatCompletionChunk(
id=f"chat{_id}",
choices=[choice],
created=_created,
model=_model,
object="chat.completion.chunk",
)
def _create_chat_completion(self, params: Dict[str, Any]) -> Union[ChatCompletion, JSONResponse]:
"""
Creates a chat completion based on the given parameters.
Args:
params (Dict[str, Any]): The parameters for generating the chat completion.
Returns:
ChatCompletion: The generated chat completion.
"""
last_output = None
for output in self._generate(params):
last_output = output
if last_output["error_code"] != 0:
return create_error_response(last_output["error_code"], last_output["text"])
function_call, finish_reason = None, "stop"
if params.get("functions") or params.get("tools"):
try:
res, function_call = self.prompt_adapter.parse_assistant_response(
last_output["text"], params.get("functions"), params.get("tools"),
)
last_output["text"] = res
except Exception as e:
traceback.print_exc()
logger.warning("Failed to parse tool call")
if isinstance(function_call, dict) and "arguments" in function_call:
finish_reason = "function_call"
function_call = FunctionCall(**function_call)
message = ChatCompletionMessage(
role="assistant",
content=last_output["text"],
function_call=function_call,
)
elif isinstance(function_call, dict) and "function" in function_call:
finish_reason = "tool_calls"
tool_calls = [model_parse(ChatCompletionMessageToolCall, function_call)]
message = ChatCompletionMessage(
role="assistant",
content=last_output["text"],
tool_calls=tool_calls,
)
else:
message = ChatCompletionMessage(
role="assistant",
content=last_output["text"].strip(),
)
choice = Choice(
index=0,
message=message,
finish_reason=finish_reason,
logprobs=None,
)
usage = model_parse(CompletionUsage, last_output["usage"])
return ChatCompletion(
id=f"chat{last_output['id']}",
choices=[choice],
created=last_output["created"],
model=last_output["model"],
object="chat.completion",
usage=usage,
)
def create_completion(
self,
params: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> Union[Iterator, Completion]:
params = params or {}
params.update(kwargs)
return (
self._create_completion_stream(params)
if params.get("stream", False)
else self._create_completion(params)
)
def create_chat_completion(
self,
params: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[Iterator, ChatCompletion]:
params = params or {}
params.update(kwargs)
return (
self._create_chat_completion_stream(params)
if params.get("stream", False)
else self._create_chat_completion(params)
)
@property
def stop(self):
"""
Gets the stop property of the prompt adapter.
Returns:
The stop property of the prompt adapter, or None if it does not exist.
"""
return self.prompt_adapter.stop if hasattr(self.prompt_adapter, "stop") else None