import copy import json import re import logging import torch from dataclasses import dataclass, field from typing import Optional, List, Dict, Any, Union, Callable, Tuple, AsyncGenerator import functools from maga_transformer.models.base_model import GenerateOutput, GenerateOutputs from maga_transformer.config.generate_config import GenerateConfig from maga_transformer.openai.renderers.qwen_tool_renderer import QwenToolRenderer, QwenToolStreamStatus from maga_transformer.tokenizer.tokenization_qwen import QWenTokenizer from transformers import Qwen2Tokenizer from maga_transformer.openai.api_datatype import ChatMessage, GPTFunctionDefinition, \ ChatCompletionRequest, RoleEnum, FunctionCall, ChatCompletionResponseStreamChoice, \ DeltaMessage, FinisheReason, UsageInfo, RendererInfo, PromptTokensDetails from maga_transformer.openai.renderers.custom_renderer import CustomChatRenderer, RendererParams, \ StreamResponseObject, RenderedInputs, StreamStatus, StreamStatusSync, OutputDelta, ThinkStatus from maga_transformer.openai.renderers.basic_renderer import BasicRenderer from maga_transformer.openai.renderer_factory_register import register_renderer from maga_transformer.utils.word_util import get_stop_word_slices, truncate_response_with_stop_words, is_truncated QwenTokenizerTypes = Union[QWenTokenizer, Qwen2Tokenizer] TOOL_DESC = """{name_for_model}: Call this tool to interact with the {name_for_human} API. What is the {name_for_human} API useful for? {description_for_model} Parameters: {parameters}""" REACT_INSTRUCTION = """Answer the following questions as best you can. You have access to the following APIs: {tools_text} Use the following format: Question: the input question you must answer Thought: you should always think about what to do Action: the action to take, should be one of [{tools_name_text}] Action Input: the input to the action Observation: the result of the action ... (this Thought/Action/Action Input/Observation can be repeated zero or more times) Thought: I now know the final answer Final Answer: the final answer to the original input question Begin!""" DUMMY_THOUGHT = { "en": "\nThought: I now know the final answer.\nFinal answer: ", "zh": "\nThought: 我会作答了。\nFinal answer: ", } _TEXT_COMPLETION_CMD = object() class QwenStreamStatus(StreamStatus): generating_function_call: bool = False total_output_string: str = "" def __init__(self, request: ChatCompletionRequest): super().__init__(request) def update_result(self): self.last_token_length = len(self.output_ids) - len(self.last_output_ids) self.last_output_ids = self.output_ids self.responded_string = self.total_output_string[: - len('\nAction:')] @property def responded_length(self): return len(self.responded_string) @property def output_length(self): return len(self.total_output_string) def check_stop_reason(self): if self.finish_reason == None: logging.debug(f"output [{self.responded_string}] found no stop reason! use stop as default.") self.finish_reason = FinisheReason.stop class QwenStreamStatusSync(StreamStatusSync): generating_function_call: bool = False total_output_string: str = "" def __init__(self, request: ChatCompletionRequest): super().__init__(request) def update_result(self): self.responded_string = self.total_output_string[: - len('\nAction:')] @property def responded_length(self): return len(self.responded_string) @property def output_length(self): return len(self.total_output_string) def check_stop_reason(self): if self.finish_reason == None: logging.debug(f"output [{self.responded_string}] found no stop reason! use stop as default.") self.finish_reason = FinisheReason.stop @dataclass class ProcessedOutput: output_str: str output_token_length: int finish_reason: Optional[FinisheReason] # TODO(wangyin): pass `max_window_size` to here. def make_context( tokenizer: QwenTokenizerTypes, query: str, history: List[Tuple[str, str]] = [], system: str = "", max_window_size: int = 6144, ): history = copy.deepcopy(history) im_start, im_end = "<|im_start|>", "<|im_end|>" im_start_tokens = [tokenizer.im_start_id] im_end_tokens = [tokenizer.im_end_id] nl_tokens = tokenizer.encode("\n") def _tokenize_str(role, content): return f"{role}\n{content}", tokenizer.encode( role, allowed_special=set() ) + nl_tokens + tokenizer.encode(content, allowed_special=set()) system_text, system_tokens_part = _tokenize_str("system", system) system_tokens = im_start_tokens + system_tokens_part + im_end_tokens raw_text = "" context_tokens = [] for turn_query, turn_response in reversed(history): query_text, query_tokens_part = _tokenize_str("user", turn_query) query_tokens = im_start_tokens + query_tokens_part + im_end_tokens response_text, response_tokens_part = _tokenize_str( "assistant", turn_response ) response_tokens = im_start_tokens + response_tokens_part + im_end_tokens next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens prev_chat = ( f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}" ) current_context_size = ( len(system_tokens) + len(next_context_tokens) + len(context_tokens) ) if current_context_size < max_window_size: context_tokens = next_context_tokens + context_tokens raw_text = prev_chat + raw_text else: break context_tokens = system_tokens + context_tokens raw_text = f"{im_start}{system_text}{im_end}" + raw_text context_tokens += ( nl_tokens + im_start_tokens + _tokenize_str("user", query)[1] + im_end_tokens + nl_tokens + im_start_tokens + tokenizer.encode("assistant") + nl_tokens ) raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n" return raw_text, context_tokens class QwenRenderer(CustomChatRenderer): def __init__(self, tokenizer: QwenTokenizerTypes, renderer_params: RendererParams): super().__init__(tokenizer, renderer_params) self.add_extra_stop_word_ids([[37763, 367, 25], [151643]]) # Observation: self.qwen_tool_renderer = QwenToolRenderer(tokenizer, renderer_params) self.template_chat_renderer: Optional[BasicRenderer] = None try: if tokenizer.chat_template != None: logging.info(f"qwen model has chat_template [{tokenizer.chat_template}], " "which will be used for non-function call dialogue.") self.template_chat_renderer = BasicRenderer(tokenizer, renderer_params) except AttributeError: pass def render_chat(self, request: ChatCompletionRequest) -> RenderedInputs: if request.tools: return self.qwen_tool_renderer.render_chat(request) if (self.template_chat_renderer != None) and \ ((request.functions == None) or (len(request.functions) == 0)): return self.template_chat_renderer.render_chat(request) query, history, system = self.parse_messages(request.messages, request.functions) logging.debug(f"parsed query: {query}, history: {history}, system: {system}") input_ids = [] if (query == _TEXT_COMPLETION_CMD): input_ids = self.text_complete_last_message(history) else: assert (isinstance(query, str)) input_ids = make_context(self.tokenizer, query, history, system)[1] return RenderedInputs(input_ids=input_ids) def text_complete_last_message(self, history): im_start = "<|im_start|>" im_end = "<|im_end|>" prompt = f"{im_start}system\nYou are a helpful assistant.{im_end}" for i, (query, response) in enumerate(history): query = query.lstrip("\n").rstrip() response = response.lstrip("\n").rstrip() prompt += f"\n{im_start}user\n{query}{im_end}" prompt += f"\n{im_start}assistant\n{response}{im_end}" prompt = prompt[: -len(im_end)] return self.tokenizer.encode(prompt) def parse_messages( self, messages: List[ChatMessage], functions: Optional[List[GPTFunctionDefinition]] = None ): if all(m.role != "user" for m in messages): raise ValueError("At least one message must be from user.") messages = copy.deepcopy(messages) if messages[0].role == 'system': system = messages.pop(0).content.lstrip('\n').rstrip() else: system = 'You are a helpful assistant.' if functions: tools_text = [] tools_name_text = [] for func_info in functions: name = func_info.name name_m = func_info.name_for_model or name name_h = func_info.name_for_human or name desc = func_info.description desc_m = func_info.description_for_model or desc tool = TOOL_DESC.format( name_for_model=name_m, name_for_human=name_h, # Hint: You can add the following format requirements in description: # "Format the arguments as a JSON object." # "Enclose the code within triple backticks (`) at the beginning and end of the code." description_for_model=desc_m, parameters=json.dumps(func_info.parameters, ensure_ascii=False), ) tools_text.append(tool) tools_name_text.append(name_m) tools_text = "\n\n".join(tools_text) tools_name_text = ", ".join(tools_name_text) instruction = (REACT_INSTRUCTION.format( tools_text=tools_text, tools_name_text=tools_name_text, ).lstrip('\n').rstrip()) else: instruction = '' messages_with_fncall = messages messages = [] for m_idx, m in enumerate(messages_with_fncall): role, content, func_call = m.role, m.content, m.function_call content = content or "" content = content.lstrip("\n").rstrip() if role == "function": if (len(messages) == 0) or (messages[-1].role != "assistant"): raise ValueError(f"Invalid request: Expecting role assistant before role function.") messages[-1].content += f'\nObservation: {content}' if m_idx == len(messages_with_fncall) - 1: # add a prefix for text completion messages[-1].content += '\nThought:' elif role == 'assistant': if len(messages) == 0: raise ValueError(f"Invalid request: Expecting role user before role assistant.") if func_call is None: if functions: content = f'Thought: I now know the final answer.\nFinal Answer: {content}' else: f_name, f_args = func_call.name, func_call.arguments if not content.startswith('Thought:'): content = f'Thought: {content}' content = f'{content}\nAction: {f_name}\nAction Input: {f_args}' if messages[-1].role == 'user': messages.append( ChatMessage(role=RoleEnum.assistant, content=content.lstrip('\n').rstrip()) ) else: messages[-1].content += '\n' + content elif role == 'user': messages.append( ChatMessage(role='user',content=content.lstrip('\n').rstrip())) else: raise ValueError(f"Invalid request: Incorrect role {role}.") query = _TEXT_COMPLETION_CMD if messages[-1].role == 'user': query = messages[-1].content messages = messages[:-1] history = [] # [(Q1, A1), (Q2, A2), ..., (Q_last_turn, A_last_turn)] for i in range(0, len(messages), 2): if messages[i].role == 'user' and messages[i + 1].role == 'assistant': usr_msg = messages[i].content.lstrip('\n').rstrip() bot_msg = messages[i + 1].content.lstrip('\n').rstrip() if instruction and (i == len(messages) - 2): usr_msg = f'{instruction}\n\nQuestion: {usr_msg}' instruction = '' history.append([usr_msg, bot_msg]) else: raise ValueError( "Invalid request: Expecting exactly one user (or function) role before every assistant role." ) if instruction: assert query is not _TEXT_COMPLETION_CMD query = f'{instruction}\n\nQuestion: {query}' return query, history, system def _parse_function_response(self, response: str) -> Optional[DeltaMessage]: func_name, func_args = "", "" i = response.rfind("\nAction:") j = response.rfind("\nAction Input:") k = response.rfind("\nObservation:") if 0 <= i < j: # If the text has `Action` and `Action input`, if k < j: # but does not contain `Observation`, # then it is likely that `Observation` is omitted by the LLM, # because the output text may have discarded the stop word. response = response.rstrip() + "\nObservation:" # Add it back. k = response.rfind("\nObservation:") func_name = response[i + len("\nAction:") : j].strip() func_args = response[j + len("\nAction Input:") : k].strip() logging.info(f"parsed function from response: [{response}]: {func_name}, {func_args}") if func_name: return DeltaMessage( content=response[:i], function_call=FunctionCall(name=func_name, arguments=func_args), ) return None # z = response.rfind("\nFinal Answer: ") # if z >= 0: # response = response[z + len("\nFinal Answer: ") :] async def _update_single_status(self, status: StreamStatus, output: GenerateOutput, max_new_tokens: int, stop_words_str: List[str], stop_word_slice_list: List[str], is_streaming: bool) -> OutputDelta: if status.request.tools: return await self.qwen_tool_renderer._update_single_status( status, output, max_new_tokens, stop_words_str, stop_word_slice_list, is_streaming) if not isinstance(status, QwenStreamStatus): return await super()._update_single_status(status, output, max_new_tokens, stop_words_str, stop_word_slice_list, is_streaming) if status.finish_reason != None: return await self._create_empty_delta(status.output.aux_info) status.update_output(output, self._clean_output_ids, functools.partial(self._check_finish_reason, max_new_tokens=max_new_tokens), self._remove_stop_word_ids) status.total_output_string = self.tokenizer.decode(status.output_ids).strip() if (len(status.total_output_string)) and (u'\uFFFD' == status.total_output_string[-1]): return await self._create_empty_delta(output.aux_info) # For some tokenizers (e.g. ChatGLM), decode a single token differs from decode a list of tokens. if (status.total_output_string.endswith("\nAction:")): status.generating_function_call = True return await self._create_empty_delta(output.aux_info) if (status.generating_function_call): return await self._create_empty_delta(output.aux_info) if is_truncated(status.total_output_string, stop_words_str, is_streaming): status.finish_reason = FinisheReason.stop return await self._create_empty_delta(output.aux_info) if (len(status.total_output_string) > status.responded_length + len('\nAction:')): status.delta_output_string = status.total_output_string[status.responded_length : status.output_length - len('\nAction:')] if is_truncated(status.delta_output_string, stop_word_slice_list, is_streaming): return await self._create_empty_delta(output.aux_info) else: status.update_result() return OutputDelta( status.delta_output_string, await self._generate_log_probs(status, output), status.input_token_length, status.output_token_length, status.reuse_length) return await self._create_empty_delta(output.aux_info) #override async def _create_status_list(self, n: int, request: ChatCompletionRequest) -> List[StreamStatus]: if request.tools: return [QwenToolStreamStatus(request) for _ in range(n)] if request.functions and (len(request.functions) > 0): return [QwenStreamStatus(request) for _ in range(n)] else: return [StreamStatus(request) for _ in range(n)] #override async def _flush_buffer(self, buffer_list: List[StreamStatus], stop_words_str: List[str], is_streaming: bool, think_status: ThinkStatus): if buffer_list[0].request.tools: return await self.qwen_tool_renderer._flush_buffer(buffer_list, stop_words_str, is_streaming, think_status) if (not isinstance(buffer_list[0], QwenStreamStatus)): return await super()._flush_buffer(buffer_list, stop_words_str, is_streaming, think_status) output_items: List[OutputDelta] = [] for status in buffer_list: if status.generating_function_call: function_message = self._parse_function_response(status.total_output_string[status.responded_length:]) if (function_message == None): logging.warning(f"output [{status.total_output_string}] failed to parse function from [{status.responded_length}]. " "regarded as normal output.") function_message = "" else: status.finish_reason = FinisheReason.function_call output_items.append(OutputDelta( function_message, await self._generate_log_probs(status, status.output), status.input_token_length, status.output_token_length, status.reuse_length)) else: trunc_string = truncate_response_with_stop_words(status.total_output_string[status.responded_length:], stop_words_str, is_streaming) output_items.append(OutputDelta( trunc_string, await self._generate_log_probs(status, status.output), status.input_token_length, status.output_token_length, status.reuse_length)) return await self._generate_stream_response(output_items, think_status) #override def _update_single_status_sync(self, status: StreamStatusSync, input_len, # output.aux_info output_len, # output.aux_info reuse_len, # output.aux_info all_probs: torch.Tensor, output_ids: torch.Tensor, max_new_tokens: int, stop_words_str: List[str], stop_word_slice_list: List[str], is_streaming: bool) -> OutputDelta: # function call is disabled when logprobs is required. if not isinstance(status, QwenStreamStatusSync): return super()._update_single_status_sync(status, input_len, output_len, reuse_len, all_probs, output_ids, max_new_tokens, stop_words_str, stop_word_slice_list, is_streaming) if status.finish_reason != None: return self._create_empty_delta_sync(input_len, output_len, reuse_len) status.update_output_sync(output_ids, input_len, self._clean_output_ids, functools.partial(self._check_finish_reason, max_new_tokens=max_new_tokens), self._remove_stop_word_ids) status.total_output_string = self.tokenizer.decode(status.output_ids).strip() if (len(status.total_output_string)) and (u'\uFFFD' == status.total_output_string[-1]): return self._create_empty_delta_sync(input_len, output_len, reuse_len) # For some tokenizers (e.g. ChatGLM), decode a single token differs from decode a list of tokens. if (status.total_output_string.endswith("\nAction:")): status.generating_function_call = True return self._create_empty_delta_sync(input_len, output_len, reuse_len) if (status.generating_function_call): return self._create_empty_delta_sync(input_len, output_len, reuse_len) if is_truncated(status.total_output_string, stop_words_str, is_streaming): status.finish_reason = FinisheReason.stop return self._create_empty_delta_sync(input_len, output_len, reuse_len) if (len(status.total_output_string) > status.responded_length + len('\nAction:')): status.delta_output_string = status.total_output_string[status.responded_length : status.output_length - len('\nAction:')] if is_truncated(status.delta_output_string, stop_word_slice_list, is_streaming): return self._create_empty_delta_sync(input_len, output_len, reuse_len) else: status.update_result() return OutputDelta( output_str=status.delta_output_string, logprobs=self._generate_log_probs_sync(status, all_probs, output_ids), input_length=input_len, output_length=output_len, reuse_length=reuse_len) return self._create_empty_delta_sync(input_len, output_len, reuse_len) #override def _create_status_list_sync(self, n: int, body: str) -> List[StreamStatusSync]: request = self.getRequest(body) if request.logprobs: return [StreamStatusSync(request) for _ in range(n)] else: return [QwenStreamStatusSync(request) for _ in range(n)] #override def _flush_buffer_sync(self, buffer_list: List[StreamStatusSync], input_len_list, output_len_list, reuse_len_list, all_probs_list, output_ids_list, stop_words_str: List[str], is_streaming: bool): if (not isinstance(buffer_list[0], QwenStreamStatusSync)): return super()._flush_buffer_sync(buffer_list, input_len_list, output_len_list, reuse_len_list, all_probs_list, output_ids_list, stop_words_str, is_streaming) output_items: List[OutputDelta] = [] for status, input_len, output_len, reuse_len, all_probs, output_ids in zip( buffer_list, input_len_list, output_len_list, reuse_len_list, all_probs_list, output_ids_list ): if status.generating_function_call: function_message = self._parse_function_response(status.total_output_string[status.responded_length:]) if (function_message == None): logging.warning(f"output [{status.total_output_string}] failed to parse function from [{status.responded_length}]. " "regarded as normal output.") function_message = "" else: status.finish_reason = FinisheReason.function_call output_items.append(OutputDelta( function_message, self._generate_log_probs_sync(status, all_probs, output_ids), input_len, output_len, reuse_len)) else: trunc_string = truncate_response_with_stop_words(status.total_output_string[status.responded_length:], stop_words_str, is_streaming) output_items.append(OutputDelta( trunc_string, self._generate_log_probs_sync(status, all_probs, output_ids), input_len, output_len, reuse_len)) return self._generate_stream_response_sync(output_items) def get_renderer_info(self) -> RendererInfo: renderer_info = super().get_renderer_info() if self.template_chat_renderer: renderer_info.template = self.template_chat_renderer.chat_template return renderer_info register_renderer('qwen', QwenRenderer) register_renderer('qwen_7b', QwenRenderer) register_renderer('qwen_13b', QwenRenderer) register_renderer('qwen_1b8', QwenRenderer) register_renderer('qwen_2', QwenRenderer) register_renderer('qwen_2_moe', QwenRenderer)