File size: 3,235 Bytes
921bf58 | 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 | """Transformers inference adapter using the unchanged native model equations."""
from dataclasses import fields
import torch
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_purrence import PurrenceConfig
from .purrence_core import ModelConfig, RecurrentLM
class PurrenceForCausalLM(PreTrainedModel, GenerationMixin):
config_class = PurrenceConfig
base_model_prefix = ""
main_input_name = "input_ids"
_supports_cache_class = False
def __init__(self, config, dtype=None):
super().__init__(config)
native = RecurrentLM(ModelConfig(**{
f.name: getattr(config, f.name) for f in fields(ModelConfig)
}))
self.embedding = native.embedding
self.block = native.block
self.queries = native.queries
def get_input_embeddings(self):
return self.embedding
def set_input_embeddings(self, value):
self.embedding = value
def get_output_embeddings(self):
return None # The native forward uses the embedding weight directly.
def read(self, sources, query):
return RecurrentLM.read(self, sources, query)
def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
return dict(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
def forward(self, input_ids=None, attention_mask=None, return_dict=None,
use_cache=False, past_key_values=None, labels=None, **kwargs):
if labels is not None or past_key_values is not None:
raise ValueError("This inference adapter does not implement training or a KV cache")
if input_ids is None or input_ids.ndim != 2 or input_ids.numel() == 0:
raise ValueError("Expected a nonempty batch of token IDs")
if attention_mask is None:
mask = torch.ones_like(input_ids, dtype=torch.bool)
else:
if attention_mask.shape != input_ids.shape:
raise ValueError("Attention mask shape differs from input IDs")
mask = attention_mask.bool()
lengths = mask.sum(-1)
if (lengths == 0).any() or (lengths > self.config.max_seq_len).any():
raise ValueError("Each sequence must contain 1 to 8192 unmasked tokens")
packed = input_ids[mask]
starts = torch.cat((lengths.new_zeros(1), lengths.cumsum(0)[:-1]))
document_starts = packed.eq(self.config.eos_token_id)
document_starts[starts] = True
offsets = document_starts.nonzero().flatten()
cu = torch.cat((offsets, offsets.new_tensor([packed.numel()]))).to(torch.int32)
positions = torch.arange(packed.numel(), device=packed.device) - torch.repeat_interleave(offsets, cu.diff().long())
with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
values = RecurrentLM.forward(self, packed, cu, positions)
logits = values.new_zeros((*input_ids.shape, self.config.vocab_size))
logits[mask] = values
if return_dict is False or (return_dict is None and not self.config.use_return_dict):
return (logits,)
return CausalLMOutput(logits=logits)
|