| """Hugging Face I/O adapter. AIMSModel owns all model computation and RoPE.""" |
|
|
| import torch |
| from torch.nn import functional as F |
| from transformers import GenerationMixin, PreTrainedModel |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
| from .configuration_baby_aims import BabyAIMSConfig |
| from .notebook_components import AIMSModel |
|
|
|
|
| class BabyAIMSForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = BabyAIMSConfig |
| base_model_prefix = "model" |
| main_input_name = "input_ids" |
| _no_split_modules = ["TransformerBlock"] |
| accepts_loss_kwargs = False |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.model = AIMSModel(config.to_notebook_config()) |
| self.post_init() |
|
|
| def _init_weights(self, module): |
| |
| |
| pass |
|
|
| @classmethod |
| def from_pretrained(cls, *args, **kwargs): |
| want_info = kwargs.pop("output_loading_info", False) |
| model, info = super().from_pretrained(*args, output_loading_info=True, **kwargs) |
| for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs"): |
| if info.get(key): |
| raise ValueError(f"Incomplete or incompatible BabyAIMS checkpoint: {key}={info[key]}") |
| |
| model.model.restore_runtime_buffers() |
| return (model, info) if want_info else model |
|
|
| def get_input_embeddings(self): |
| return self.model.tok_emb |
|
|
| def set_input_embeddings(self, value): |
| self.model.tok_emb = value |
|
|
| def get_output_embeddings(self): |
| return self.model.out_head |
|
|
| def set_output_embeddings(self, value): |
| self.model.out_head = value |
|
|
| @staticmethod |
| def _validate_positions(input_ids, attention_mask, position_ids): |
| batch, length = input_ids.shape |
| if position_ids is not None: |
| canonical = torch.arange(length, device=input_ids.device).expand(batch, -1) |
| if attention_mask is not None: |
| canonical = (attention_mask.long().cumsum(-1) - 1).clamp_min(0) |
| if position_ids.shape not in ((1, length), (batch, length)): |
| raise ValueError("position_ids must have shape [1 or batch, sequence_length].") |
| positions = position_ids.to(input_ids.device).expand(batch, -1) |
| valid = torch.ones_like(input_ids, dtype=torch.bool) if attention_mask is None else attention_mask |
| if not torch.equal(positions[valid], canonical[valid]): |
| raise ValueError("Custom position_ids are unsupported; AIMS uses sequential positions.") |
|
|
|
|
| def forward( |
| self, input_ids=None, attention_mask=None, position_ids=None, |
| inputs_embeds=None, labels=None, past_key_values=None, use_cache=None, |
| output_attentions=None, output_hidden_states=None, return_dict=None, |
| cache_position=None, logits_to_keep=None, num_items_in_batch=None, **kwargs, |
| ): |
| if past_key_values is not None or use_cache: |
| raise ValueError("KV caching is not implemented; use use_cache=False.") |
| if inputs_embeds is not None or input_ids is None: |
| raise ValueError("Supply input_ids; inputs_embeds is not supported by AIMSModel.") |
| if output_attentions or (output_attentions is None and self.config.output_attentions): |
| raise ValueError("Returning attention weights is not implemented.") |
| if output_hidden_states or (output_hidden_states is None and self.config.output_hidden_states): |
| raise ValueError("Returning hidden states is not implemented.") |
| |
| if kwargs: |
| raise TypeError(f"Unsupported forward arguments: {sorted(kwargs)}") |
| if input_ids.ndim != 2 or input_ids.shape[0] == 0: |
| raise ValueError("input_ids must have nonempty shape [batch, sequence_length].") |
| batch, length = input_ids.shape |
| if length < 1 or length > self.config.context_length: |
| raise ValueError(f"Sequence length must be between 1 and {self.config.context_length}.") |
| if attention_mask is not None: |
| if attention_mask.shape != input_ids.shape: |
| raise ValueError("attention_mask must have shape [batch, sequence_length].") |
| attention_mask = attention_mask.to(input_ids.device) |
| if not torch.all((attention_mask == 0) | (attention_mask == 1)): |
| raise ValueError("attention_mask must contain only 0 and 1.") |
| attention_mask = attention_mask.bool() |
| if not attention_mask.any(dim=-1).all(): |
| raise ValueError("Each sequence must contain at least one unmasked token.") |
| |
| starts = attention_mask[:, :1].long().sum(-1) + ( |
| ~attention_mask[:, :-1] & attention_mask[:, 1:] |
| ).sum(-1) |
| if (starts > 1).any(): |
| raise ValueError("attention_mask must not contain internal gaps.") |
| self._validate_positions(input_ids, attention_mask, position_ids) |
|
|
| if attention_mask is None or bool(attention_mask.all()): |
| logits = self.model(input_ids) |
| else: |
| rows = [] |
| for row in range(batch): |
| indices = attention_mask[row].nonzero(as_tuple=True)[0] |
| compact_ids = input_ids[row:row + 1].index_select(1, indices) |
| compact_logits = self.model(compact_ids) |
| restored = compact_logits.new_zeros(length, self.config.vocab_size) |
| rows.append(restored.index_copy(0, indices, compact_logits[0])) |
| logits = torch.stack(rows) |
|
|
| |
| |
| if not torch.isfinite(logits).all(): |
| raise ValueError("AIMSModel produced NaN/Inf logits before sampling; check weights, dtype and runtime buffers.") |
| loss = None |
| if labels is not None: |
| if labels.shape != input_ids.shape: |
| raise ValueError("labels must match input shape and must not be pre-shifted.") |
| shifted_labels = labels[:, 1:].to(logits.device).clone() |
| if attention_mask is not None: |
| valid_pairs = attention_mask[:, :-1] & attention_mask[:, 1:] |
| shifted_labels.masked_fill_(~valid_pairs, -100) |
| if (shifted_labels != -100).any(): |
| loss = F.cross_entropy( |
| logits[:, :-1].contiguous().float().view(-1, self.config.vocab_size), |
| shifted_labels.reshape(-1), ignore_index=-100, |
| ) |
| else: |
| loss = logits.sum() * 0.0 |
| return_dict = self.config.return_dict if return_dict is None else return_dict |
| if not return_dict: |
| return (loss, logits) if loss is not None else (logits,) |
| return CausalLMOutputWithPast(loss=loss, logits=logits) |
|
|
| def prepare_inputs_for_generation( |
| self, input_ids, attention_mask=None, inputs_embeds=None, |
| past_key_values=None, use_cache=False, **kwargs, |
| ): |
| if inputs_embeds is not None or past_key_values is not None or use_cache: |
| raise ValueError("Generation supports input_ids and use_cache=False only.") |
| if attention_mask is not None and not attention_mask[:, -1].bool().all(): |
| raise ValueError("Use left padding for batched generation.") |
| |
| result = {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False} |
| if kwargs.get("position_ids") is not None: |
| |
| |
| positions = kwargs["position_ids"] |
| prefix_length = positions.shape[-1] |
| if prefix_length > input_ids.shape[1]: |
| raise ValueError("position_ids cannot exceed the input sequence length.") |
| prefix_mask = None if attention_mask is None else attention_mask[:, :prefix_length].bool() |
| self._validate_positions(input_ids[:, :prefix_length], prefix_mask, positions) |
| for name in ("output_attentions", "output_hidden_states"): |
| if name in kwargs: |
| result[name] = kwargs[name] |
| return result |
|
|
|
|
| BabyAIMSForCausalLM.register_for_auto_class("AutoModelForCausalLM") |
|
|
