106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
|
# Copyright (c) 2022 Zhipu.AI
|
||
|
|
import copy
|
||
|
|
from typing import Optional, Tuple, Union, List, Callable, Dict, Any
|
||
|
|
from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig
|
||
|
|
from transformers.generation.logits_process import LogitsProcessor
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from modelscope.models.base import Model, TorchModel
|
||
|
|
from modelscope.models.builder import MODELS
|
||
|
|
from modelscope.outputs import OutputKeys
|
||
|
|
from modelscope.utils.constant import Tasks
|
||
|
|
from modelscope.utils.logger import get_logger
|
||
|
|
from modelscope.pipelines.base import Pipeline
|
||
|
|
from modelscope.pipelines.builder import PIPELINES
|
||
|
|
from modelscope.preprocessors import Preprocessor
|
||
|
|
from modelscope.utils.constant import Tasks
|
||
|
|
|
||
|
|
|
||
|
|
from transformers import AutoTokenizer,LlamaForCausalLM
|
||
|
|
from transformers import PreTrainedModel, PreTrainedTokenizer
|
||
|
|
|
||
|
|
|
||
|
|
@MODELS.register_module(Tasks.text_generation, module_name='minigpt7b')
|
||
|
|
class MiniGPT7bForTextGeneration(TorchModel):
|
||
|
|
|
||
|
|
def __init__(self, model_dir: str, *args, **kwargs):
|
||
|
|
"""initialize the minigpt7b from the `model_dir` path.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model_dir (str): the model path.
|
||
|
|
"""
|
||
|
|
super().__init__(model_dir, *args, **kwargs)
|
||
|
|
self.logger = get_logger()
|
||
|
|
self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
|
||
|
|
# loading tokenizer
|
||
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
||
|
|
# loading model
|
||
|
|
self.model = LlamaForCausalLM.from_pretrained(model_dir,torch_dtype=torch.float16)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def answer(self, inputs, max_new_tokens=300, num_beams=1, min_length=1, top_p=0.9,
|
||
|
|
repetition_penalty=1.0, length_penalty=1, temperature=1.0, max_length=2000,device=None):
|
||
|
|
|
||
|
|
input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids
|
||
|
|
input_ids = input_ids.to(device)
|
||
|
|
model_ids = self.model.to(device)
|
||
|
|
outputs = model_ids.generate(
|
||
|
|
#inputs_embeds=embs,
|
||
|
|
input_ids,
|
||
|
|
max_new_tokens=max_new_tokens,
|
||
|
|
#stopping_criteria=self.stopping_criteria,
|
||
|
|
num_beams=num_beams,
|
||
|
|
do_sample=True,
|
||
|
|
min_length=min_length,
|
||
|
|
top_p=top_p,
|
||
|
|
repetition_penalty=repetition_penalty,
|
||
|
|
length_penalty=length_penalty,
|
||
|
|
temperature=temperature,
|
||
|
|
)
|
||
|
|
|
||
|
|
rets = self.tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
||
|
|
|
||
|
|
return rets
|
||
|
|
|
||
|
|
|
||
|
|
def forward(self, input: Dict) -> Dict:
|
||
|
|
output = {}
|
||
|
|
res = self.answer(input,device=self.device)
|
||
|
|
output['text'] = res
|
||
|
|
return output
|
||
|
|
|
||
|
|
def quantize(self, bits: int):
|
||
|
|
self.model = self.model.quantize(bits)
|
||
|
|
return self
|
||
|
|
|
||
|
|
|
||
|
|
@PIPELINES.register_module(
|
||
|
|
group_key=Tasks.text_generation,
|
||
|
|
module_name='minigpt-7b-text-generation')
|
||
|
|
class MiniGPT7bTextGenerationPipeline(Pipeline):
|
||
|
|
|
||
|
|
def __init__(self,
|
||
|
|
model: Union[Model, str],
|
||
|
|
preprocessor: [Preprocessor] = None,
|
||
|
|
*args,
|
||
|
|
**kwargs):
|
||
|
|
model = MiniGPT7bForTextGeneration(model) if isinstance(model,
|
||
|
|
str) else model
|
||
|
|
self.model = model
|
||
|
|
self.model.eval()
|
||
|
|
|
||
|
|
super().__init__(model=model, **kwargs)
|
||
|
|
|
||
|
|
def preprocess(self, inputs, **preprocess_params) -> Dict[str, Any]:
|
||
|
|
return inputs
|
||
|
|
|
||
|
|
# define the forward pass
|
||
|
|
def forward(self, inputs: Dict, **forward_params) -> Dict[str, Any]:
|
||
|
|
return self.model(inputs)
|
||
|
|
|
||
|
|
# format the outputs from pipeline
|
||
|
|
def postprocess(self, input, **kwargs) -> Dict[str, Any]:
|
||
|
|
return input
|
||
|
|
|