from transformers import PreTrainedTokenizer, GenerationConfig, StoppingCriteriaList from typing import Optional, Callable, List, Tuple, Union import copy import torch from transformers import AutoTokenizer from transformers.generation.logits_process import LogitsProcessorList from packaging import version _ERROR_BAD_CHAT_FORMAT = """\ We detect you are probably using the pretrained model (rather than chat model) for chatting, since the chat_format in generation_config is not "chatml". If you are directly using the model downloaded from Huggingface, please make sure you are using our "Qwen/Qwen-7B-Chat" Huggingface model (rather than "Qwen/Qwen-7B") when you call model.chat(). 我们检测到您可能在使用预训练模型（而非chat模型）进行多轮chat，因为您当前在generation_config指定的chat_format，并未设置为我们在对话中所支持的"chatml"格式。 如果您在直接使用我们从Huggingface提供的模型，请确保您在调用model.chat()时，使用的是"Qwen/Qwen-7B-Chat"模型（而非"Qwen/Qwen-7B"预训练模型）。 """ IMEND = "<|im_end|>" ENDOFTEXT = "<|endoftext|>" HistoryType = List[Tuple[str, str]] TokensType = List[int] BatchTokensType = List[List[int]] def get_stop_words_ids(chat_format, tokenizer): if chat_format == "raw": stop_words_ids = [tokenizer.encode("Human:"), [tokenizer.eod_id]] elif chat_format == "chatml": stop_words_ids = [[tokenizer.im_end_id], [tokenizer.im_start_id]] else: raise NotImplementedError(f"Unknown chat format {chat_format!r}") return stop_words_ids def make_context( tokenizer: PreTrainedTokenizer, query: str, history: List[Tuple[str, str]] = None, system: str = "", max_window_size: int = 6144, chat_format: str = "chatml", ): if history is None: history = [] if chat_format == "chatml": 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" elif chat_format == "raw": raw_text = query context_tokens = tokenizer.encode(raw_text) else: raise NotImplementedError(f"Unknown chat format {chat_format!r}") return raw_text, context_tokens class vLLMWrapper: def __init__(self, model_dir: str, trust_remote_code: bool = True, tensor_parallel_size: int = 1, gpu_memory_utilization: float = 0.98, dtype: str = "bfloat16", **kwargs): if dtype not in ("bfloat16", "float16", "float32"): print("now not support {}!".format(dtype)) raise Exception # build generation_config self.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=trust_remote_code) # build tokenizer self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) self.tokenizer.eos_token_id = self.generation_config.eos_token_id self.stop_words_ids = [] from vllm import LLM import vllm if version.parse(vllm.__version__) >= version.parse("0.2.2"): self.__vllm_support_repetition_penalty = True else: self.__vllm_support_repetition_penalty = False quantization = getattr(kwargs, 'quantization', None) self.model = LLM(model=model_dir, tokenizer=model_dir, tensor_parallel_size=tensor_parallel_size, trust_remote_code=trust_remote_code, quantization=quantization, gpu_memory_utilization=gpu_memory_utilization, dtype=dtype) for stop_id in get_stop_words_ids(self.generation_config.chat_format, self.tokenizer): self.stop_words_ids.extend(stop_id) self.stop_words_ids.extend([self.generation_config.eos_token_id]) def chat(self, query: str, history: Optional[HistoryType], tokenizer: PreTrainedTokenizer = None, system: str = "You are a helpful assistant.", generation_config: Optional[GenerationConfig] = None, **kwargs): generation_config = generation_config if generation_config is not None else self.generation_config tokenizer = self.tokenizer if tokenizer is None else tokenizer assert generation_config.chat_format == 'chatml', _ERROR_BAD_CHAT_FORMAT if not self.__vllm_support_repetition_penalty and generation_config.repetition_penalty != 1: raise RuntimeError("The installed vLLM doesn't support repetition_penalty, please set ``model.generation_config.repetition_penalty = 1`` or install vllm>=0.2.2") if history is None: history = [] else: # make a copy of the user's input such that is is left untouched history = copy.deepcopy(history) extra_stop_words_ids = kwargs.get('stop_words_ids', None) if extra_stop_words_ids is None: extra_stop_words_ids = [] max_window_size = kwargs.get('max_window_size', None) if max_window_size is None: max_window_size = generation_config.max_window_size from vllm.sampling_params import SamplingParams sampling_kwargs = { "stop_token_ids": self.stop_words_ids, "early_stopping": False, "top_p": generation_config.top_p, "top_k": -1 if generation_config.top_k == 0 else generation_config.top_k, "temperature": generation_config.temperature, "max_tokens": generation_config.max_new_tokens, "repetition_penalty": generation_config.repetition_penalty } if not self.__vllm_support_repetition_penalty: sampling_kwargs.pop("repetition_penalty") sampling_params = SamplingParams(**sampling_kwargs) raw_text, context_tokens = make_context( self.tokenizer, query, history=history, system=system, max_window_size=max_window_size, chat_format=generation_config.chat_format, ) req_outputs = self.model.generate([query], sampling_params=sampling_params, prompt_token_ids=[context_tokens]) req_output = req_outputs[0] prompt_str = req_output.prompt prompt_ids = req_output.prompt_token_ids req_sample_output_ids = [] req_sample_output_strs = [] for sample in req_output.outputs: output_str = sample.text output_ids = sample.token_ids if IMEND in output_str: output_str = output_str[:-len(IMEND)] if ENDOFTEXT in output_str: output_str = output_str[:-len(ENDOFTEXT)] req_sample_output_ids.append(prompt_ids + output_ids) req_sample_output_strs.append(prompt_str + output_str) assert len(req_sample_output_strs) == 1 response = req_sample_output_strs[0][len(prompt_str):] history.append((prompt_str, response)) return response, history if __name__ == '__main__': model_dir = 'Qwen/Qwen-72B-Chat' tensor_parallel_size = 2 model = vLLMWrapper(model_dir, tensor_parallel_size=tensor_parallel_size, ) response, history = model.chat(query="你好", history=None) print(response) response, history = model.chat(query="给我讲一个年轻人奋斗创业最终取得成功的故事。", history=history) print(response) response, history = model.chat(query="给这个故事起一个标题", history=history) print(response)