137 lines
5.3 KiB
Python
137 lines
5.3 KiB
Python
|
import math
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
from transformers import GPT2PreTrainedModel
|
||
|
|
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
|
||
|
|
|
||
|
|
|
||
|
|
class GPT2InferenceModel(GPT2PreTrainedModel):
|
||
|
|
"""Override GPT2LMHeadModel to allow for prefix conditioning."""
|
||
|
|
|
||
|
|
def __init__(self, config, gpt, pos_emb, embeddings, norm, linear, kv_cache):
|
||
|
|
super().__init__(config)
|
||
|
|
self.transformer = gpt
|
||
|
|
self.pos_embedding = pos_emb
|
||
|
|
self.embeddings = embeddings
|
||
|
|
self.final_norm = norm
|
||
|
|
self.lm_head = nn.Sequential(norm, linear)
|
||
|
|
self.kv_cache = kv_cache
|
||
|
|
|
||
|
|
def store_prefix_emb(self, prefix_emb):
|
||
|
|
self.cached_prefix_emb = prefix_emb
|
||
|
|
|
||
|
|
def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
|
||
|
|
token_type_ids = kwargs.get("token_type_ids", None) # usually None
|
||
|
|
if not self.kv_cache:
|
||
|
|
past_key_values = None
|
||
|
|
|
||
|
|
# only last token for inputs_ids if past is defined in kwargs
|
||
|
|
if past_key_values is not None:
|
||
|
|
input_ids = input_ids[:, -1].unsqueeze(-1)
|
||
|
|
if token_type_ids is not None:
|
||
|
|
token_type_ids = token_type_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_key_values is not None:
|
||
|
|
position_ids = position_ids[:, -1].unsqueeze(-1)
|
||
|
|
else:
|
||
|
|
position_ids = None
|
||
|
|
return {
|
||
|
|
"input_ids": input_ids,
|
||
|
|
"past_key_values": past_key_values,
|
||
|
|
"use_cache": kwargs.get("use_cache"),
|
||
|
|
"position_ids": position_ids,
|
||
|
|
"attention_mask": attention_mask,
|
||
|
|
"token_type_ids": token_type_ids,
|
||
|
|
}
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
input_ids=None,
|
||
|
|
past_key_values=None,
|
||
|
|
attention_mask=None,
|
||
|
|
token_type_ids=None,
|
||
|
|
position_ids=None,
|
||
|
|
head_mask=None,
|
||
|
|
inputs_embeds=None,
|
||
|
|
encoder_hidden_states=None,
|
||
|
|
encoder_attention_mask=None,
|
||
|
|
labels=None,
|
||
|
|
use_cache=None,
|
||
|
|
output_attentions=None,
|
||
|
|
output_hidden_states=None,
|
||
|
|
return_dict=None,
|
||
|
|
):
|
||
|
|
assert self.cached_prefix_emb is not None
|
||
|
|
assert inputs_embeds is None # Not supported by this inference model.
|
||
|
|
assert labels is None # Training not supported by this inference model.
|
||
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||
|
|
|
||
|
|
# assert len(past_key_values) + len(input_ids) == attention_mask.shape[1]
|
||
|
|
|
||
|
|
# Create embedding
|
||
|
|
prefix_len = self.cached_prefix_emb.shape[1]
|
||
|
|
if input_ids.shape[1] != 1:
|
||
|
|
gen_inputs = input_ids[:, prefix_len:]
|
||
|
|
gen_emb = self.embeddings(gen_inputs)
|
||
|
|
gen_emb = gen_emb + self.pos_embedding(gen_emb)
|
||
|
|
if self.cached_prefix_emb.shape[0] != gen_emb.shape[0]:
|
||
|
|
prefix_emb = self.cached_prefix_emb.repeat_interleave(
|
||
|
|
gen_emb.shape[0] // self.cached_prefix_emb.shape[0], 0
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
prefix_emb = self.cached_prefix_emb.to(gen_emb.dtype)
|
||
|
|
emb = torch.cat([prefix_emb, gen_emb], dim=1)
|
||
|
|
else:
|
||
|
|
emb = self.embeddings(input_ids)
|
||
|
|
emb = emb + self.pos_embedding.get_fixed_embedding(
|
||
|
|
attention_mask.shape[1] - (prefix_len + 1), attention_mask.device
|
||
|
|
)
|
||
|
|
transformer_outputs = self.transformer(
|
||
|
|
inputs_embeds=emb,
|
||
|
|
past_key_values=past_key_values,
|
||
|
|
attention_mask=attention_mask,
|
||
|
|
token_type_ids=token_type_ids,
|
||
|
|
position_ids=position_ids,
|
||
|
|
head_mask=head_mask,
|
||
|
|
encoder_hidden_states=encoder_hidden_states,
|
||
|
|
encoder_attention_mask=encoder_attention_mask,
|
||
|
|
use_cache=use_cache,
|
||
|
|
output_attentions=output_attentions,
|
||
|
|
output_hidden_states=output_hidden_states,
|
||
|
|
return_dict=return_dict,
|
||
|
|
)
|
||
|
|
hidden_states = transformer_outputs[0]
|
||
|
|
lm_logits = self.lm_head(hidden_states)
|
||
|
|
|
||
|
|
if not return_dict:
|
||
|
|
return (lm_logits,) + transformer_outputs[1:]
|
||
|
|
|
||
|
|
return CausalLMOutputWithCrossAttentions(
|
||
|
|
loss=None,
|
||
|
|
logits=lm_logits,
|
||
|
|
past_key_values=transformer_outputs.past_key_values,
|
||
|
|
hidden_states=transformer_outputs.hidden_states,
|
||
|
|
attentions=transformer_outputs.attentions,
|
||
|
|
cross_attentions=transformer_outputs.cross_attentions,
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _reorder_cache(past, beam_idx):
|
||
|
|
"""
|
||
|
|
This function is used to re-order the :obj:`past_key_values` cache if
|
||
|
|
:meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
|
||
|
|
called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
|
||
|
|
"""
|
||
|
|
return tuple(
|
||
|
|
tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
|
||
|
|
for layer_past in past
|
||
|
|
)
|