99 lines
3.7 KiB
Python
99 lines
3.7 KiB
Python
# Copyright (c) 2022 Zhipu.AI
|
|
import torch
|
|
from modelscope.pipelines.builder import PIPELINES
|
|
from modelscope.pipelines.base import Pipeline
|
|
from modelscope.utils.constant import Tasks
|
|
from modelscope.models.builder import MODELS
|
|
from modelscope.utils.logger import get_logger
|
|
from modelscope.models.base import Model, TorchModel
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from transformers import PreTrainedModel, PreTrainedTokenizer
|
|
import numpy as np
|
|
from typing import Union, Dict, Any
|
|
|
|
PROMPT_FORMAT = """Below is an instruction that describes a task. Write a response that appropriately completes the request.
|
|
|
|
### Instruction:
|
|
{instruction}
|
|
|
|
### Response:
|
|
"""
|
|
|
|
def generate_response(instruction: str, *, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
|
|
do_sample: bool = True, max_new_tokens: int = 256, top_p: float = 0.92, top_k: int = 0, device=None, **kwargs) -> str:
|
|
input_ids = tokenizer(PROMPT_FORMAT.format(instruction=instruction), return_tensors="pt").input_ids
|
|
# each of these is encoded to a single token
|
|
response_key_token_id = tokenizer.encode("### Response:")[0]
|
|
end_key_token_id = tokenizer.encode("### End")[0]
|
|
input_ids = input_ids.to(torch.device(device))
|
|
model = model.to(torch.device(device))
|
|
gen_tokens = model.generate(input_ids, pad_token_id=tokenizer.pad_token_id, eos_token_id=end_key_token_id,
|
|
do_sample=do_sample, max_new_tokens=max_new_tokens, top_p=top_p, top_k=top_k, **kwargs)[0].cpu()
|
|
|
|
# find where the response begins
|
|
response_positions = np.where(gen_tokens == response_key_token_id)[0]
|
|
|
|
if len(response_positions) >= 0:
|
|
response_pos = response_positions[0]
|
|
|
|
# find where the response ends
|
|
end_pos = None
|
|
end_positions = np.where(gen_tokens == end_key_token_id)[0]
|
|
if len(end_positions) > 0:
|
|
end_pos = end_positions[0]
|
|
|
|
return tokenizer.decode(gen_tokens[response_pos + 1 : end_pos]).strip()
|
|
|
|
return None
|
|
|
|
|
|
@PIPELINES.register_module(Tasks.text_generation, module_name='DollyV16b-text-generation')
|
|
class DollyV16bTextGenerationPipeline(Pipeline):
|
|
def __init__(self,
|
|
model: Union[Model, str],
|
|
*args,
|
|
**kwargs):
|
|
device = kwargs.get('device')
|
|
model = DollyV16bTextGeneration(model_dir=model,device=device) 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
|
|
|
|
|
|
@MODELS.register_module(Tasks.text_generation, module_name='DollyV16b')
|
|
class DollyV16bTextGeneration(TorchModel):
|
|
|
|
def __init__(self, model_dir: str, device=None, *args, **kwargs):
|
|
super().__init__(model_dir, *args, **kwargs)
|
|
self.logger = get_logger()
|
|
self.device = device
|
|
# loading tokenizer
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir, padding_side="left")
|
|
# loading model
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True)
|
|
|
|
|
|
|
|
|
|
def forward(self, input: Dict) -> Dict:
|
|
output = {}
|
|
res = generate_response(input,model=self.model,tokenizer=self.tokenizer,device=self.device)
|
|
output['text'] = res
|
|
return output
|
|
|
|
def quantize(self, bits: int):
|
|
self.model = self.model.quantize(bits)
|
|
return self
|
|
|