diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 3d9ba089..ab22b9ba 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -523,6 +523,7 @@ class Pipelines(object): soonet_video_temporal_grounding = 'soonet-video-temporal-grounding' efficient_diffusion_tuning = 'efficient-diffusion-tuning' multimodal_dialogue = 'multimodal-dialogue' + llama2_text_generation_pipeline = 'llama2-text-generation-pipeline' # science tasks protein_structure = 'unifold-protein-structure' diff --git a/modelscope/pipelines/nlp/llama2_text_generation_pipeline.py b/modelscope/pipelines/nlp/llama2_text_generation_pipeline.py new file mode 100644 index 00000000..3a9d3d44 --- /dev/null +++ b/modelscope/pipelines/nlp/llama2_text_generation_pipeline.py @@ -0,0 +1,99 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +# Copyright (c) 2022 Zhipu.AI +from typing import Any, Dict, Union + +import torch + +from modelscope import Model, snapshot_download +from modelscope.metainfo import Pipelines, Preprocessors +from modelscope.models.nlp.llama2 import Llama2Tokenizer +from modelscope.pipelines.base import Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.nlp.text_generation_pipeline import \ + TextGenerationPipeline +from modelscope.preprocessors import Preprocessor +from modelscope.utils.constant import Fields, Tasks + + +@PIPELINES.register_module( + Tasks.text_generation, + module_name=Pipelines.llama2_text_generation_pipeline) +class Llama2TaskPipeline(TextGenerationPipeline): + + def __init__(self, + model: Union[Model, str], + preprocessor: Preprocessor = None, + config_file: str = None, + device: str = 'gpu', + auto_collate=True, + **kwargs): + """Use `model` and `preprocessor` to create a generation pipeline for prediction. + + Args: + model (str or Model): Supply either a local model dir which supported the text generation task, + or a model id from the model hub, or a torch model instance. + preprocessor (Preprocessor): An optional preprocessor instance, please make sure the preprocessor fits for + the model if supplied. + kwargs (dict, `optional`): + Extra kwargs passed into the preprocessor's constructor. + Examples: + >>> from modelscope.utils.constant import Tasks + >>> import torch + >>> from modelscope.pipelines import pipeline + >>> from modelscope import snapshot_download, Model + >>> model_dir = snapshot_download("modelscope/Llama-2-13b-chat-ms", + >>> ignore_file_pattern = [r'\\w+\\.safetensors']) + >>> pipe = pipeline(task=Tasks.text_generation, model=model_dir, device_map='auto', + >>> torch_dtype=torch.float16) + >>> inputs="咖啡的作用是什么？" + >>> result = pipe(inputs,max_length=200, do_sample=True, top_p=0.85, + >>> temperature=1.0, repetition_penalty=1., eos_token_id=2, bos_token_id=1, pad_token_id=0) + >>> print(result['text']) + + To view other examples plese check tests/pipelines/test_llama2_text_generation_pipeline.py. + """ + self.model = Model.from_pretrained( + model, device_map='auto', torch_dtype=torch.float16) + self.tokenizer = Llama2Tokenizer.from_pretrained(model) + super().__init__(model=self.model, **kwargs) + + def preprocess(self, inputs, **preprocess_params) -> Dict[str, Any]: + return inputs + + def _sanitize_parameters(self, **pipeline_parameters): + return {}, pipeline_parameters, {} + + def forward(self, + inputs, + max_length=50, + do_sample=True, + top_p=0.85, + temperature=1.0, + repetition_penalty=1., + eos_token_id=2, + bos_token_id=1, + pad_token_id=0, + **forward_params) -> Dict[str, Any]: + output = {} + inputs = self.tokenizer(inputs, return_tensors='pt') + generate_ids = self.model.generate( + inputs.input_ids.to('cuda'), + max_length=max_length, + do_sample=do_sample, + top_p=top_p, + temperature=temperature, + repetition_penalty=repetition_penalty, + eos_token_id=eos_token_id, + bos_token_id=bos_token_id, + pad_token_id=pad_token_id, + **forward_params) + out = self.tokenizer.batch_decode( + generate_ids, + skip_special_tokens=True, + clean_up_tokenization_spaces=False)[0] + output['text'] = out + return output + + # format the outputs from pipeline + def postprocess(self, input, **kwargs) -> Dict[str, Any]: + return input diff --git a/tests/pipelines/test_llama2_text_generation_pipeline.py b/tests/pipelines/test_llama2_text_generation_pipeline.py new file mode 100644 index 00000000..2a532257 --- /dev/null +++ b/tests/pipelines/test_llama2_text_generation_pipeline.py @@ -0,0 +1,47 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import unittest + +import torch + +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.test_utils import test_level + + +class Llama2TextGenerationPipelineTest(unittest.TestCase): + + def setUp(self) -> None: + self.llama2_model_id_7B_chat_ms = 'modelscope/Llama-2-7b-chat-ms' + self.llama2_input_chat_ch = '天空为什么是蓝色的？' + + def run_pipeline_with_model_id(self, + model_id, + input, + init_kwargs={}, + run_kwargs={}): + pipeline_ins = pipeline( + task=Tasks.text_generation, model=model_id, **init_kwargs) + pipeline_ins._model_prepare = True + result = pipeline_ins(input, **run_kwargs) + print(result['text']) + + # 7B_ms_chat + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_llama2_7B_chat_ms_with_model_name_with_chat_ch_with_args(self): + self.run_pipeline_with_model_id( + self.llama2_model_id_7B_chat_ms, + self.llama2_input_chat_ch, + init_kwargs={ + 'device_map': 'auto', + 'torch_dtype': torch.float16 + }, + run_kwargs={ + 'max_length': 200, + 'do_sample': True, + 'top_p': 0.85 + }) + + +if __name__ == '__main__': + unittest.main()