From a0fe5947b7b4f9f10286681437005463df580d47 Mon Sep 17 00:00:00 2001 From: Firmament-cyou <57580313+Firmament-cyou@users.noreply.github.com> Date: Tue, 5 Mar 2024 10:05:18 +0800 Subject: [PATCH] Support stream_generate for LLMPipeline (#768) * support streaming output for llm_pipeline * add qwen2 format_messages --- modelscope/pipelines/nlp/llm_pipeline.py | 90 ++++++++++++++++++++---- tests/pipelines/test_llm_pipeline.py | 15 ++++ 2 files changed, 93 insertions(+), 12 deletions(-) diff --git a/modelscope/pipelines/nlp/llm_pipeline.py b/modelscope/pipelines/nlp/llm_pipeline.py index 3f641f76..1d2effd0 100644 --- a/modelscope/pipelines/nlp/llm_pipeline.py +++ b/modelscope/pipelines/nlp/llm_pipeline.py @@ -1,12 +1,11 @@ # Copyright (c) Alibaba, Inc. and its affiliates. import os -import os.path as osp from contextlib import contextmanager -from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Generator, Iterator, List, Tuple, Union import json import torch -from transformers import PreTrainedTokenizer +from transformers import PreTrainedModel, PreTrainedTokenizer from modelscope import (AutoModelForCausalLM, AutoTokenizer, Pipeline, snapshot_download) @@ -14,12 +13,17 @@ from modelscope.hub.file_download import model_file_download from modelscope.models.base import Model from modelscope.models.nlp import ChatGLM2Tokenizer, Llama2Tokenizer from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input from modelscope.pipelines.builder import PIPELINES from modelscope.pipelines.util import is_model, is_official_hub_path from modelscope.utils.config import Config -from modelscope.utils.constant import Invoke, ModelFile, Tasks +from modelscope.utils.constant import Frameworks, Invoke, ModelFile, Tasks +from modelscope.utils.device import device_placement from modelscope.utils.logger import get_logger from modelscope.utils.model_type_helper import ModelTypeHelper +from modelscope.utils.streaming_output import (PipelineStreamingOutputMixin, + StreamingOutputMixin, + add_stream_generate) logger = get_logger() @@ -72,7 +76,7 @@ class LLMAdapterRegistry: @PIPELINES.register_module(Tasks.chat, module_name='llm') @PIPELINES.register_module(Tasks.text_generation, module_name='llm') -class LLMPipeline(Pipeline): +class LLMPipeline(Pipeline, PipelineStreamingOutputMixin): def initiate_single_model(self, model): if isinstance(model, str): @@ -168,6 +172,8 @@ class LLMPipeline(Pipeline): self.ignore_file_pattern = kwargs.pop('ignore_file_pattern', None) with self._temp_configuration_file(kwargs): super().__init__(*args, **kwargs) + if isinstance(self.model, PreTrainedModel): + self.model = add_stream_generate(self.model) tokenizer_class = None if isinstance(format_messages, str): @@ -207,8 +213,9 @@ class LLMPipeline(Pipeline): forward_params = kwargs.get('forward_params', {}) postprocess_params = kwargs.get('postprocess_params', {}) - is_messages = isinstance(inputs, dict) and 'messages' in inputs - tokens = self.preprocess(inputs, is_messages, **preprocess_params) + preprocess_params['is_messages'] = postprocess_params['is_messages'] \ + = isinstance(inputs, dict) and 'messages' in inputs + tokens = self.preprocess(inputs, **preprocess_params) if self.llm_framework is None: # pytorch model @@ -226,11 +233,62 @@ class LLMPipeline(Pipeline): if self.llm_framework is None: # pytorch model outputs = outputs.tolist()[0][len(tokens['inputs'][0]):] - response = self.postprocess(outputs, is_messages, **postprocess_params) + response = self.postprocess(outputs, **postprocess_params) return response - def preprocess(self, inputs: Union[str, Dict], is_messages: bool, - **kwargs): + def stream_generate(self, inputs: Union[Input, List[Input]], *args, + **kwargs) -> Generator: + assert isinstance(self.model, StreamingOutputMixin + ), 'pipeline.model must be StreamingOutputMixin!' + if (self.model or (self.has_multiple_models and self.models[0])): + if not self._model_prepare: + self.prepare_model() + + preprocess_params, forward_params, postprocess_params = self._sanitize_parameters( + **kwargs) + preprocess_params['is_messages'] = postprocess_params['is_messages'] \ + = isinstance(inputs, dict) and 'messages' in inputs + + if isinstance(inputs, list): + model_input_list = [ + self._preprocess_with_check(i, preprocess_params) + for i in inputs + ] + output = [] + for ele in model_input_list: + output.append( + self._stream_single(ele, forward_params, + postprocess_params)) + else: + model_input = self._preprocess_with_check(inputs, + preprocess_params) + output = self._stream_single(model_input, forward_params, + postprocess_params) + return output + + def _stream_single(self, model_input: Dict[str, Any], + forward_params: Dict[str, Any], + postprocess_params: Dict[str, Any]) -> Generator: + + with device_placement(self.framework, self.device_name): + if self.framework == Frameworks.torch: + with torch.no_grad(): + if self._auto_collate: + model_input = self._collate_fn(model_input) + stream = self.model.stream_generate( + **model_input, **forward_params) + else: + stream = self.model.stream_generate(**model_input, + **forward_params) + + for out in stream: + out = out.tolist()[0][len(model_input['inputs'][0]):] + out = self.postprocess(out, **postprocess_params) + self._check_output(out) + yield out + + def preprocess(self, inputs: Union[str, Dict], **kwargs): + is_messages = kwargs.pop('is_messages') if is_messages: tokens = self.format_messages(inputs, self.tokenizer, **kwargs) else: @@ -252,8 +310,8 @@ class LLMPipeline(Pipeline): for k, v in tokens.items() } - def postprocess(self, outputs, is_messages: bool, **kwargs): - + def postprocess(self, outputs, **kwargs): + is_messages = kwargs.pop('is_messages') if not isinstance(outputs, str): response = self.tokenizer.decode( outputs, skip_special_tokens=True, **kwargs) @@ -569,6 +627,14 @@ def chatglm3_format_messages(messages, tokenizer, **kwargs): return inputs +@LLMAdapterRegistry.register_format_messages('qwen2') +def qwen2_format_messages(messages, tokenizer, **kwargs): + messages = messages['messages'] + text = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True) + return tokenizer([text], return_tensors='pt') + + LLMAdapterRegistry.register_tokenizer('chatglm2', ChatGLM2Tokenizer) LLMAdapterRegistry.register_tokenizer('llama', Llama2Tokenizer) LLMAdapterRegistry.register_tokenizer('llama2', Llama2Tokenizer) diff --git a/tests/pipelines/test_llm_pipeline.py b/tests/pipelines/test_llm_pipeline.py index 47653071..94c2f168 100644 --- a/tests/pipelines/test_llm_pipeline.py +++ b/tests/pipelines/test_llm_pipeline.py @@ -350,6 +350,21 @@ class LLMPipelineTest(unittest.TestCase): print('messages: ', pipe(self.messages_zh, **self.gen_cfg)) print('prompt: ', pipe(self.prompt_zh, **self.gen_cfg)) + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_qwen_stream_gemerate(self): + pipe = pipeline(task='chat', model='qwen/Qwen-7B-Chat', llm_first=True) + for stream_output in pipe.stream_generate(self.messages_zh_with_system, + **self.gen_cfg): + print('messages: ', stream_output, end='\r') + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_qwen1_5_stream_gemerate(self): + pipe = pipeline( + task='chat', model='qwen/Qwen1.5-1.8B-Chat', llm_first=True) + for stream_output in pipe.stream_generate(self.messages_zh_with_system, + **self.gen_cfg): + print('messages: ', stream_output, end='\r') + if __name__ == '__main__': unittest.main()