226 lines
8.6 KiB
Python
226 lines
8.6 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
from typing import Any, Dict
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers.modeling_utils import PreTrainedModel
|
|
|
|
from modelscope.metainfo import TaskModels
|
|
from modelscope.models.builder import MODELS
|
|
from modelscope.models.nlp.task_models.task_model import \
|
|
SingleBackboneTaskModelBase
|
|
from modelscope.outputs import (OutputKeys, TextGenerationModelOutput,
|
|
TokenGeneratorOutput)
|
|
from modelscope.utils.constant import Tasks
|
|
|
|
__all__ = ['ModelForTextGeneration']
|
|
|
|
|
|
@MODELS.register_module(
|
|
Tasks.text_generation, module_name=TaskModels.text_generation)
|
|
class ModelForTextGeneration(SingleBackboneTaskModelBase, PreTrainedModel):
|
|
|
|
def __init__(self, model_dir: str, *args, **kwargs):
|
|
"""initialize the text generation model from the `model_dir` path.
|
|
|
|
Args:
|
|
model_dir (str): the model path.
|
|
"""
|
|
super().__init__(model_dir, *args, **kwargs)
|
|
if 'base_model_prefix' in kwargs:
|
|
self._base_model_prefix = kwargs['base_model_prefix']
|
|
|
|
self.build_backbone(self.backbone_cfg)
|
|
self.build_head(self.head_cfg)
|
|
if self.config.get('shared_embedding', False):
|
|
input_embeddings = self.backbone.get_input_embeddings()
|
|
output_embeddings = self.head.get_output_embeddings()
|
|
output_embeddings.weight = input_embeddings.weight
|
|
|
|
def forward(self, **inputs) -> Dict[str, np.ndarray]:
|
|
# backbone do not need labels, only head need for loss compute
|
|
labels = inputs.pop(OutputKeys.LABELS, None)
|
|
|
|
backbone_outputs = super().forward(inputs)
|
|
hidden_states = backbone_outputs[0]
|
|
|
|
logits = self.head.forward(hidden_states)
|
|
loss = None
|
|
if labels is not None:
|
|
inputs[OutputKeys.LABELS] = labels
|
|
loss = self.compute_loss(logits, labels)
|
|
return TextGenerationModelOutput(logits=logits, loss=loss)
|
|
|
|
def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
|
|
# only last token for inputs_ids if past is defined in kwargs
|
|
if past:
|
|
input_ids = input_ids[:, -1].unsqueeze(-1)
|
|
|
|
attention_mask = kwargs.get('attention_mask', None)
|
|
position_ids = kwargs.get('position_ids', None)
|
|
|
|
if attention_mask is not None and position_ids is None:
|
|
# create position_ids on the fly for batch generation
|
|
position_ids = attention_mask.long().cumsum(-1) - 1
|
|
position_ids.masked_fill_(attention_mask == 0, 1)
|
|
if past:
|
|
position_ids = position_ids[:, -1].unsqueeze(-1)
|
|
else:
|
|
position_ids = None
|
|
return {
|
|
'input_ids': input_ids,
|
|
'past_key_values': past,
|
|
'use_cache': kwargs.get('use_cache'),
|
|
'position_ids': position_ids,
|
|
'attention_mask': attention_mask,
|
|
}
|
|
|
|
def generate(self, inputs, temperature=1.0, **kwargs):
|
|
tokens = inputs['input_ids'] if isinstance(inputs, Dict) else inputs
|
|
top_k = kwargs.pop('top_k',
|
|
self.config.top_k if 'top_k' in self.config else 1)
|
|
top_p = kwargs.pop('top_p',
|
|
self.config.top_p if 'top_p' in self.config else 0.)
|
|
max_length = kwargs.pop('max_length', self.config.max_length)
|
|
|
|
batch_size = tokens.size(0)
|
|
lengths = kwargs.pop(
|
|
'prompt_length',
|
|
torch.tensor([tokens.size(1)], device=tokens.device))
|
|
|
|
min_prompt_length = lengths.min().item()
|
|
max_sequence_length = max_length
|
|
|
|
# If the context is too big, this happens
|
|
if min_prompt_length >= max_sequence_length:
|
|
raise ValueError('context length too large')
|
|
|
|
pad_length = max_sequence_length - tokens.size(1)
|
|
if pad_length > 0:
|
|
pads = torch.zeros(
|
|
batch_size, pad_length, device=tokens.device).long()
|
|
tokens = torch.cat((tokens, pads), dim=-1)
|
|
|
|
# Added termination_id to support the case that we want to terminate the
|
|
# generation once that id is generated.
|
|
termination_id = self.config.eos_token_id
|
|
|
|
# Whether we have reached a termination id.
|
|
is_generation_done = torch.zeros(
|
|
batch_size, dtype=torch.uint8, device=tokens.device)
|
|
|
|
with torch.no_grad():
|
|
for context_length in range(min_prompt_length,
|
|
max_sequence_length):
|
|
|
|
# Pick the slice that we need to pass through the network.
|
|
tokens2use = tokens[:, :context_length]
|
|
|
|
# logits will be meanigful only in the last pipeline stage.
|
|
logits = self(input_ids=tokens2use).logits
|
|
|
|
# Sample.
|
|
last_token_logits = logits[:, -1, :]
|
|
new_sample = sample(
|
|
last_token_logits,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
temperature=temperature,
|
|
vocab_size=self.head_cfg.vocab_size)
|
|
|
|
# If a prompt length is smaller or equal th current context
|
|
# length, it means we have started generating tokens
|
|
started = lengths <= context_length
|
|
# Update the tokens.
|
|
tokens[started, context_length] = new_sample[started]
|
|
|
|
done_token = (new_sample == termination_id).byte() & \
|
|
started.byte()
|
|
|
|
is_generation_done = is_generation_done | done_token
|
|
done = torch.all(is_generation_done)
|
|
|
|
if done:
|
|
break
|
|
|
|
tokens = tokens[:, :(context_length + 1)]
|
|
return TokenGeneratorOutput(sequences=tokens)
|
|
|
|
|
|
def sample(logits, top_k=0, top_p=0.0, temperature=1.0, vocab_size=None):
|
|
""" Sample and generate a token.
|
|
Note: logits has the dimension [b, v] where b is the batch size
|
|
and v is the vocabulary size.
|
|
If vocab_size is provided, we will make sure the sample that is
|
|
generated is in [0, vocab-size). This will avoid out of vocabulary
|
|
generations due to padding.
|
|
"""
|
|
|
|
# Check logits for consistency.
|
|
assert logits.ndim == 2, 'expected the logits to be of [b, v] shape.'
|
|
|
|
# Greedy is just simple argmax.
|
|
if top_k == 1:
|
|
assert top_p == 0.0, 'cannot set both greedy and top-p samplings.'
|
|
samples = torch.argmax(logits, dim=-1)
|
|
|
|
# Top-k or top-p sampling.
|
|
else:
|
|
# Clone so we do not modify the inputs,
|
|
logits = logits.clone()
|
|
# Apply temperature in place.
|
|
if temperature != 1.0:
|
|
logits.div_(temperature)
|
|
|
|
if top_k > 1:
|
|
top_p == 0.0
|
|
assert top_k <= logits.size(1), 'top-k is larger than logit size.'
|
|
if vocab_size:
|
|
assert top_k < vocab_size, 'top-k is larger than vocab size.'
|
|
modify_logits_for_top_k_filtering(logits, top_k)
|
|
|
|
elif top_p > 0.0:
|
|
assert top_p <= 1.0, 'top-p should be in (0, 1].'
|
|
modify_logits_for_top_p_filtering(logits, top_p)
|
|
|
|
# After filtering, we need to recalculate the distribution.
|
|
probs = logits.softmax(dim=-1)
|
|
samples = torch.multinomial(probs, num_samples=1).view(-1)
|
|
|
|
# If vocab size is provided, make sure the samples are in
|
|
# in the range [0, vocab-size).
|
|
if vocab_size:
|
|
samples = torch.clamp(samples, min=0, max=(vocab_size - 1))
|
|
|
|
return samples
|
|
|
|
|
|
def modify_logits_for_top_k_filtering(logits, top_k):
|
|
"""Set the logits for none top-k values to -inf."""
|
|
|
|
filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
|
|
logits.masked_fill_(filter_, float('-Inf'))
|
|
|
|
|
|
def modify_logits_for_top_p_filtering(logits, top_p):
|
|
"""Set the logits for none top-p values to -inf."""
|
|
|
|
# First sort and calculate cumulative sum of probabilities.
|
|
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
|
cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
|
|
|
|
# Filteration based on the cumulative sum.
|
|
filter_ = cumulative_probs > top_p
|
|
# This shift by 1 is weird and I cannot justify it. This existed
|
|
# in the original implementation:
|
|
# https://github.com/ari-holtzman/degen/blob/master/gen.py
|
|
# and I guess it is needed so keeping it for now.
|
|
filter_[:, 1:] = filter_[:, :-1].clone()
|
|
# Make sure we at least have one token to select from.
|
|
filter_[..., 0] = 0
|
|
|
|
# Fill in the filtered part
|
|
filter_ = filter_.scatter(1, sorted_indices, filter_)
|
|
logits.masked_fill_(filter_, float('-Inf'))
|