File size: 3,305 Bytes
a22f353 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import torch
from torch.nn import functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_sol_pro import SolProHFConfig
from .modeling_sol_pro import SolProConfig, SolProForCausalLM
class SolProHFForCausalLM(PreTrainedModel, GenerationMixin):
config_class = SolProHFConfig
base_model_prefix = ""
main_input_name = "input_ids"
_no_split_modules = ["Block"]
def __init__(self, config):
super().__init__(config)
fields = SolProConfig.__dataclass_fields__
core = SolProForCausalLM(SolProConfig(**{k: getattr(config, k) for k in fields}))
# Keep original checkpoint names, including the single tied embedding.
for name, module in core.named_children():
self.add_module(name, module)
for name, parameter in core.named_parameters(recurse=False):
self.register_parameter(name, parameter)
for name, buffer in core.named_buffers(recurse=False):
self.register_buffer(name, buffer)
self.post_init()
def get_input_embeddings(self):
return self.embedding
def set_input_embeddings(self, value):
self.embedding = value
def get_output_embeddings(self):
return self.embedding
def set_output_embeddings(self, value):
self.embedding = value
def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}
def forward(self, input_ids=None, attention_mask=None, labels=None,
return_dict=None, use_cache=False, past_key_values=None, **kwargs):
if input_ids is None:
raise ValueError("input_ids are required")
if past_key_values is not None:
raise ValueError("KV caching is not implemented")
if attention_mask is None or bool(attention_mask.bool().all()):
logits = SolProForCausalLM.forward(self, input_ids)
else:
# Strip padding so n-gram memory and recurrent conditioning see the
# same sequence as an unpadded standalone inference call.
rows = []
for ids, mask in zip(input_ids, attention_mask):
keep = mask.bool()
positions = keep.nonzero(as_tuple=True)[0]
if positions.numel() == 0:
raise ValueError("An attention mask cannot be empty")
values = SolProForCausalLM.forward(self, ids[keep][None])[0]
row = values.new_zeros((ids.numel(), self.config.vocab_size))
rows.append(row.index_copy(0, positions, values))
logits = torch.stack(rows)
loss = None
if labels is not None:
targets = labels[:, 1:].clone()
if attention_mask is not None:
targets.masked_fill_(~attention_mask[:, 1:].bool(), -100)
loss = F.cross_entropy(logits[:, :-1].float().reshape(-1, logits.shape[-1]),
targets.reshape(-1), ignore_index=-100)
if return_dict is False:
return ((loss,) if loss is not None else ()) + (logits,)
return CausalLMOutput(loss=loss, logits=logits)
|