| """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 |
|
|
| 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) |
|
|
