"""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): # Original PyTorch constructors own initialization. Incomplete pretrained # checkpoints are rejected below, never filled with uninitialized storage. 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]}") # Public loading boundary; the original model owns its buffer lifecycle. 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.") # Recognized HF bookkeeping above never changes original computation. 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.") # A single contiguous run supports left/right padding, not interior holes. 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) # Check raw logits, before HF temperature/top-k/top-p can transform them. # Never silently sanitize corrupted model outputs. 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.") # Preserve flags for validation; bookkeeping never slices the prefix. result = {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False} if kwargs.get("position_ids") is not None: # HF may retain the caller's original prompt positions across steps. # Validate that prefix once per step, then let AIMS own all positions. 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")