Spaces:
Running
Running
| import json | |
| import re | |
| from copy import deepcopy | |
| from typing import List, Union, Optional, Dict, Any, Tuple | |
| from fastapi import HTTPException | |
| from loguru import logger | |
| from openai.types.chat import ( | |
| ChatCompletionMessageParam, | |
| ChatCompletionUserMessageParam, | |
| ChatCompletionAssistantMessageParam, | |
| ) | |
| from transformers import PreTrainedTokenizer | |
| from api.generation.utils import parse_messages | |
| from api.utils.protocol import Role | |
| 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!""" | |
| _TEXT_COMPLETION_CMD = object() | |
| def build_qwen_chat_input( | |
| tokenizer: PreTrainedTokenizer, | |
| messages: List[ChatCompletionMessageParam], | |
| context_len: int = 8192, | |
| max_new_tokens: int = 256, | |
| functions: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, | |
| tools: Optional[List[Dict[str, Any]]] = None, | |
| ) -> List[int]: | |
| """ | |
| Builds the input tokens for Qwen chat generation. | |
| Refs: | |
| https://huggingface.co/Qwen/Qwen-7B-Chat/blob/main/qwen_generation_utils.py | |
| Args: | |
| tokenizer: The tokenizer used to encode the input tokens. | |
| messages: The list of chat messages. | |
| context_len: The maximum length of the context. | |
| max_new_tokens: The maximum number of new tokens to add. | |
| functions: Optional dictionary or list of dictionaries representing the functions. | |
| tools: Optional list of dictionaries representing the tools. | |
| Returns: | |
| The list of input tokens. | |
| """ | |
| query, history = process_qwen_messages(messages, functions, tools) | |
| if query is _TEXT_COMPLETION_CMD: | |
| return build_last_message_input(tokenizer, history) | |
| messages = [] | |
| for q, r in history: | |
| messages.extend( | |
| [ | |
| ChatCompletionUserMessageParam(role="user", content=q), | |
| ChatCompletionAssistantMessageParam(role="assistant", content=r) | |
| ] | |
| ) | |
| messages.append(ChatCompletionUserMessageParam(role="user", content=query)) | |
| max_input_tokens = context_len - max_new_tokens | |
| system, rounds = parse_messages(messages) | |
| system = f"You are a helpful assistant.{system}" | |
| im_start_tokens, im_end_tokens = [tokenizer.im_start_id], [tokenizer.im_end_id] | |
| nl_tokens = tokenizer.encode("\n") | |
| def _tokenize_str(role, content): | |
| return tokenizer.encode( | |
| role, allowed_special=set() | |
| ) + nl_tokens + tokenizer.encode(content, allowed_special=set()) | |
| system_tokens_part = _tokenize_str("system", system) | |
| system_tokens = im_start_tokens + system_tokens_part + im_end_tokens | |
| max_history_tokens = max_input_tokens - len(system_tokens) | |
| history_tokens = [] | |
| for r in rounds[::-1]: | |
| round_tokens = [] | |
| for message in r: | |
| if round_tokens: | |
| round_tokens += nl_tokens | |
| if message["role"] == Role.USER: | |
| content_tokens = im_start_tokens + _tokenize_str("user", message["content"]) + im_end_tokens | |
| else: | |
| content_tokens = im_start_tokens + _tokenize_str("assistant", message["content"]) + im_end_tokens | |
| round_tokens.extend(content_tokens) | |
| if len(history_tokens) == 0 or len(history_tokens) + len(round_tokens) <= max_history_tokens: | |
| if history_tokens: | |
| history_tokens = nl_tokens + history_tokens | |
| history_tokens = round_tokens + history_tokens # concat left | |
| if len(history_tokens) < max_history_tokens: | |
| continue | |
| break | |
| input_tokens = system_tokens + nl_tokens + history_tokens | |
| if messages[-1]["role"] != Role.ASSISTANT: | |
| input_tokens += nl_tokens + im_start_tokens + tokenizer.encode("assistant") + nl_tokens | |
| return input_tokens[-max_input_tokens:] # truncate left | |
| def check_is_qwen(model) -> bool: | |
| """ | |
| Checks if the given model is a Qwen model. | |
| Args: | |
| model: The model to be checked. | |
| Returns: | |
| bool: True if the model is a Qwen model, False otherwise. | |
| """ | |
| return "QWenBlock" in getattr(model, "_no_split_modules", []) | |
| def process_qwen_messages( | |
| messages: List[ChatCompletionMessageParam], | |
| functions: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, | |
| tools: Optional[List[Dict[str, Any]]] = None, | |
| ) -> Tuple[str, List[List[str]]]: | |
| """ | |
| Process the Qwen messages and generate a query and history. | |
| Args: | |
| messages (List[ChatCompletionMessageParam]): The list of chat completion messages. | |
| functions (Optional[Union[Dict[str, Any], List[Dict[str, Any]]]]): The functions to be used. | |
| tools (Optional[List[Dict[str, Any]]]): The tools to be used. | |
| Returns: | |
| Tuple[str, List[List[str]]]: The generated query and history. | |
| """ | |
| if all(m["role"] != Role.USER for m in messages): | |
| raise HTTPException( | |
| status_code=400, | |
| detail=f"Invalid request: Expecting at least one user message.", | |
| ) | |
| messages = deepcopy(messages) | |
| default_system = "You are a helpful assistant." | |
| system = "" | |
| if messages[0]["role"] == Role.SYSTEM: | |
| system = messages.pop(0)["content"].lstrip("\n").rstrip() | |
| if system == default_system: | |
| system = "" | |
| if tools: | |
| functions = [t["function"] for t in tools] | |
| if functions: | |
| tools_text = [] | |
| tools_name_text = [] | |
| for func_info in functions: | |
| name = func_info.get("name", "") | |
| name_m = func_info.get("name_for_model", name) | |
| name_h = func_info.get("name_for_human", name) | |
| desc = func_info.get("description", "") | |
| desc_m = func_info.get("description_for_model", 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) | |
| system += "\n\n" + REACT_INSTRUCTION.format( | |
| tools_text=tools_text, | |
| tools_name_text=tools_name_text, | |
| ) | |
| system = system.lstrip("\n").rstrip() | |
| dummy_thought = { | |
| "en": "\nThought: I now know the final answer.\nFinal answer: ", | |
| "zh": "\nThought: 我会作答了。\nFinal answer: ", | |
| } | |
| _messages = messages | |
| messages = [] | |
| for m_idx, m in enumerate(_messages): | |
| role, content = m["role"], m["content"] | |
| func_call, tool_calls = m.get("function_call", None), m.get("tool_calls", None) | |
| if content: | |
| content = content.lstrip("\n").rstrip() | |
| if role in [Role.FUNCTION, Role.TOOL]: | |
| if (len(messages) == 0) or (messages[-1]["role"] != Role.ASSISTANT): | |
| raise HTTPException( | |
| status_code=400, | |
| detail=f"Invalid request: Expecting role assistant before role function.", | |
| ) | |
| messages[-1]["content"] += f"\nObservation: {content}" | |
| if m_idx == len(_messages) - 1: | |
| messages[-1]["content"] += "\nThought:" | |
| elif role == Role.ASSISTANT: | |
| if len(messages) == 0: | |
| raise HTTPException( | |
| status_code=400, | |
| detail=f"Invalid request: Expecting role user before role assistant.", | |
| ) | |
| last_msg = messages[-1]["content"] | |
| last_msg_has_zh = len(re.findall(r"[\u4e00-\u9fff]+", last_msg)) > 0 | |
| if func_call is None and tool_calls is None: | |
| if functions or tool_calls: | |
| content = dummy_thought["zh" if last_msg_has_zh else "en"] + content | |
| else: | |
| if func_call: | |
| f_name, f_args = func_call.get("name"), func_call.get("arguments") | |
| else: | |
| f_name, f_args = tool_calls[0]["function"]["name"], tool_calls[0]["function"]["arguments"] | |
| if not content: | |
| if last_msg_has_zh: | |
| content = f"Thought: 我可以使用 {f_name} API。" | |
| else: | |
| content = f"Thought: I can use {f_name}." | |
| if messages[-1]["role"] == Role.USER: | |
| messages.append( | |
| ChatCompletionAssistantMessageParam(role="assistant", content=content.lstrip("\n").rstrip()) | |
| ) | |
| else: | |
| messages[-1]["content"] += content | |
| elif role == Role.USER: | |
| messages.append( | |
| ChatCompletionUserMessageParam(role="user", content=content.lstrip("\n").rstrip()) | |
| ) | |
| else: | |
| raise HTTPException( | |
| status_code=400, detail=f"Invalid request: Incorrect role {role}." | |
| ) | |
| query = _TEXT_COMPLETION_CMD | |
| if messages[-1]["role"] == Role.USER: | |
| query = messages[-1]["content"] | |
| messages = messages[:-1] | |
| if len(messages) % 2 != 0: | |
| raise HTTPException(status_code=400, detail="Invalid request") | |
| history = [] # [(Q1, A1), (Q2, A2), ..., (Q_last_turn, A_last_turn)] | |
| for i in range(0, len(messages), 2): | |
| if messages[i]["role"] == Role.USER and messages[i + 1]["role"] == Role.ASSISTANT: | |
| usr_msg = messages[i]["content"].lstrip("\n").rstrip() | |
| bot_msg = messages[i + 1]["content"].lstrip("\n").rstrip() | |
| if system and (i == len(messages) - 2): | |
| usr_msg = f"{system}\n\nQuestion: {usr_msg}" | |
| system = "" | |
| for t in dummy_thought.values(): | |
| t = t.lstrip("\n") | |
| if bot_msg.startswith(t) and ("\nAction: " in bot_msg): | |
| bot_msg = bot_msg[len(t):] | |
| history.append([usr_msg, bot_msg]) | |
| else: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="Invalid request: Expecting exactly one user (or function) role before every assistant role.", | |
| ) | |
| if system: | |
| assert query is not _TEXT_COMPLETION_CMD | |
| query = f"{system}\n\nQuestion: {query}" | |
| return query, history | |
| def build_last_message_input(tokenizer: PreTrainedTokenizer, history: list): | |
| 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)] | |
| logger.debug(f"==== Prompt with tools ====\n{prompt}") | |
| return tokenizer.encode(prompt) | |