""" hf_wrapper.py — Self-contained HuggingFace wrapper for BabyLMModel. All model code is inlined here so trust_remote_code=True works without needing sibling files (config.py / model.py) to be separately downloaded. Usage: from transformers import AutoConfig, AutoModelForCausalLM cfg = AutoConfig.from_pretrained("pakphum/babylm2026-ipa", trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained("pakphum/babylm2026-ipa", trust_remote_code=True) """ from __future__ import annotations import math from dataclasses import dataclass, field from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import CausalLMOutput # ───────────────────────────────────────────────────────────────────────────── # ModelConfig (inlined from config.py) # ───────────────────────────────────────────────────────────────────────────── @dataclass class ModelConfig: vocab_size: int = 32000 hidden_size: int = 512 num_layers: int = 6 num_heads: int = 8 ffn_intermediate_size: int = 1380 max_seq_len: int = 512 dropout: float = 0.1 ipa_dim: int = 24 ipa_lambda: float = 1.0 ipa_mask_ratio: float = 0.15 variant: str = "ipa_full" learning_rate: float = 3e-4 warmup_steps: int = 1000 weight_decay: float = 0.1 batch_size: int = 32 grad_clip: float = 1.0 languages: list = field(default_factory=lambda: ["en", "nl", "mandarin"]) byte_premiums: dict = field(default_factory=lambda: { "en": 1.0, "nl": 1.0516, "mandarin": 0.9894 }) word_budget: int = 100_000_000 data_dir: str = "/N/slate/partkaew/BigRed200/babylm2026/data" output_dir: str = "checkpoints" checkpoint_milestones: list = field(default_factory=lambda: [ *range(1_000_000, 10_000_000, 1_000_000), *range(10_000_000, 100_000_000, 10_000_000), *range(100_000_000, 1_000_000_000, 100_000_000), ]) def validate(self): assert self.variant in {"baseline", "ipa_add", "ipa_gate", "ipa_full"} assert self.hidden_size % self.num_heads == 0 return self # ───────────────────────────────────────────────────────────────────────────── # Model components (inlined from model.py) # ───────────────────────────────────────────────────────────────────────────── @dataclass class ModelOutput: logits: torch.Tensor aux_loss: torch.Tensor class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6) -> None: super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: rms_inv = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() return self.weight * (x * rms_inv) def _precompute_freqs_cis(head_dim: int, max_seq_len: int, theta: float = 10_000.0) -> torch.Tensor: freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) t = torch.arange(max_seq_len, dtype=torch.float32) freqs = torch.outer(t, freqs) return torch.polar(torch.ones_like(freqs), freqs) def _apply_rotary_emb(q, k, freqs_cis): T = q.shape[1] def rotate(x): x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) f = freqs_cis[:T].unsqueeze(0).unsqueeze(2) return torch.view_as_real(x_c * f).flatten(3).type_as(x) return rotate(q), rotate(k) class SwiGLUFFN(nn.Module): def __init__(self, hidden_size: int, intermediate_size: int) -> None: super().__init__() self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) def _build_attention_bias(T, attn_mask, device, dtype): causal = torch.zeros(T, T, device=device, dtype=dtype) causal = causal.masked_fill( torch.ones(T, T, device=device, dtype=torch.bool).triu(diagonal=1), float("-inf") ) bias = causal.unsqueeze(0).unsqueeze(0) if attn_mask is not None: pad = (attn_mask == 0).unsqueeze(1).unsqueeze(2) pad_bias = torch.zeros_like(pad, dtype=dtype).masked_fill(pad, float("-inf")) bias = bias + pad_bias return bias class CausalSelfAttention(nn.Module): def __init__(self, cfg: ModelConfig) -> None: super().__init__() self.n_heads = cfg.num_heads self.head_dim = cfg.hidden_size // cfg.num_heads self.attn_drop = cfg.dropout self.q_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False) self.k_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False) self.v_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False) self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False) def forward(self, x, freqs_cis, attn_mask=None): B, T, _ = x.shape q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim) k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim) v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim) q, k = _apply_rotary_emb(q, k, freqs_cis) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) bias = _build_attention_bias(T, attn_mask, x.device, x.dtype) out = F.scaled_dot_product_attention( q, k, v, attn_mask=bias, dropout_p=self.attn_drop if self.training else 0.0, is_causal=False, ) out = out.transpose(1, 2).contiguous().view(B, T, -1) return self.o_proj(out) class TransformerBlock(nn.Module): def __init__(self, cfg: ModelConfig) -> None: super().__init__() self.attn_norm = RMSNorm(cfg.hidden_size) self.attn = CausalSelfAttention(cfg) self.ffn_norm = RMSNorm(cfg.hidden_size) self.ffn = SwiGLUFFN(cfg.hidden_size, cfg.ffn_intermediate_size) def forward(self, x, freqs_cis, attn_mask=None): x = x + self.attn(self.attn_norm(x), freqs_cis, attn_mask) x = x + self.ffn(self.ffn_norm(x)) return x class IPAFusion(nn.Module): def __init__(self, cfg: ModelConfig) -> None: super().__init__() self.variant = cfg.variant self.mask_ratio = cfg.ipa_mask_ratio self.ipa_proj = nn.Linear(cfg.ipa_dim, cfg.hidden_size, bias=False) if cfg.variant in ("ipa_gate", "ipa_full"): self.gate_proj = nn.Linear(cfg.ipa_dim, cfg.hidden_size, bias=False) def forward(self, embed, ipa): ipa_target = ipa if self.training and self.variant == "ipa_full" and self.mask_ratio > 0: keep_prob = 1.0 - self.mask_ratio mask = torch.bernoulli( torch.full(ipa.shape[:2], keep_prob, device=ipa.device) ).unsqueeze(-1) ipa = ipa * mask ipa_emb = self.ipa_proj(ipa) if self.variant == "ipa_add": return embed + ipa_emb, ipa_target gate = torch.sigmoid(self.gate_proj(ipa)) return embed + gate * ipa_emb, ipa_target class IPAAuxHead(nn.Module): def __init__(self, hidden_size: int, ipa_dim: int) -> None: super().__init__() self.proj = nn.Linear(hidden_size, ipa_dim, bias=False) def loss(self, hidden, ipa_target): pred = self.proj(hidden) has_ipa = (ipa_target.abs().sum(dim=-1) > 0).float() n_ipa = has_ipa.sum() if n_ipa == 0: return hidden.new_zeros(()) mse = F.mse_loss(pred, ipa_target, reduction="none").mean(dim=-1) return (mse * has_ipa).sum() / n_ipa class BabyLMModel(nn.Module): def __init__(self, cfg: ModelConfig) -> None: super().__init__() cfg.validate() self.cfg = cfg self.embed = nn.Embedding(cfg.vocab_size, cfg.hidden_size) self.ipa_fusion = None if cfg.variant == "baseline" else IPAFusion(cfg) self.emb_drop = nn.Dropout(cfg.dropout) self.layers = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.num_layers)]) self.norm = RMSNorm(cfg.hidden_size) self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) self.lm_head.weight = self.embed.weight self.ipa_aux_head = IPAAuxHead(cfg.hidden_size, cfg.ipa_dim) if cfg.variant == "ipa_full" else None freqs_cis = _precompute_freqs_cis(cfg.hidden_size // cfg.num_heads, cfg.max_seq_len) self.register_buffer("freqs_cis", freqs_cis, persistent=False) self.apply(self._init_weights) scale = (2.0 * cfg.num_layers) ** -0.5 for name, p in self.named_parameters(): if name.endswith(("o_proj.weight", "down_proj.weight")): nn.init.normal_(p, mean=0.0, std=0.02 * scale) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, input_ids, ipa_vectors, attention_mask=None): B, T = input_ids.shape x = self.embed(input_ids) ipa_target = ipa_vectors if self.ipa_fusion is not None: x, ipa_target = self.ipa_fusion(x, ipa_vectors) x = self.emb_drop(x) freqs_cis = self.freqs_cis for layer in self.layers: x = layer(x, freqs_cis, attention_mask) x = self.norm(x) logits = self.lm_head(x) if self.ipa_aux_head is not None: aux_loss = self.ipa_aux_head.loss(x, ipa_target) else: aux_loss = logits.new_zeros(()) return ModelOutput(logits=logits, aux_loss=aux_loss) # ───────────────────────────────────────────────────────────────────────────── # HuggingFace Config # ───────────────────────────────────────────────────────────────────────────── class BabyLMConfig(PretrainedConfig): model_type = "babylm_ipa" def __init__( self, vocab_size: int = 32000, hidden_size: int = 512, num_layers: int = 6, num_heads: int = 8, ffn_intermediate_size: int = 1380, max_seq_len: int = 512, dropout: float = 0.1, ipa_dim: int = 24, ipa_lambda: float = 0.1, ipa_mask_ratio: float = 0.15, variant: str = "ipa_full", **kwargs, ): super().__init__(**kwargs) self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_layers = num_layers self.num_heads = num_heads self.ffn_intermediate_size = ffn_intermediate_size self.max_seq_len = max_seq_len self.dropout = dropout self.ipa_dim = ipa_dim self.ipa_lambda = ipa_lambda self.ipa_mask_ratio = ipa_mask_ratio self.variant = variant @classmethod def from_model_config(cls, cfg: ModelConfig) -> "BabyLMConfig": return cls( vocab_size = cfg.vocab_size, hidden_size = cfg.hidden_size, num_layers = cfg.num_layers, num_heads = cfg.num_heads, ffn_intermediate_size = cfg.ffn_intermediate_size, max_seq_len = cfg.max_seq_len, dropout = cfg.dropout, ipa_dim = cfg.ipa_dim, ipa_lambda = cfg.ipa_lambda, ipa_mask_ratio = cfg.ipa_mask_ratio, variant = cfg.variant, ) def to_model_config(self) -> ModelConfig: return ModelConfig( vocab_size = self.vocab_size, hidden_size = self.hidden_size, num_layers = self.num_layers, num_heads = self.num_heads, ffn_intermediate_size = self.ffn_intermediate_size, max_seq_len = self.max_seq_len, dropout = self.dropout, ipa_dim = self.ipa_dim, ipa_lambda = self.ipa_lambda, ipa_mask_ratio = self.ipa_mask_ratio, variant = self.variant, ) # ───────────────────────────────────────────────────────────────────────────── # HuggingFace Model # ───────────────────────────────────────────────────────────────────────────── class BabyLMForCausalLM(PreTrainedModel): """HuggingFace wrapper around BabyLMModel. The model accepts an optional `ipa_vectors` kwarg ([B, T, ipa_dim]). When omitted, zeros are used — equivalent to the baseline variant. """ config_class = BabyLMConfig supports_gradient_checkpointing = False _tied_weights_keys = {} def __init__(self, config: BabyLMConfig) -> None: super().__init__(config) model_cfg = config.to_model_config() self.model = BabyLMModel(model_cfg) self.post_init() def tie_weights(self, **kwargs) -> None: self.model.lm_head.weight = self.model.embed.weight def get_input_embeddings(self) -> nn.Embedding: return self.model.embed def set_input_embeddings(self, value: nn.Embedding) -> None: self.model.embed = value def get_output_embeddings(self) -> nn.Linear: return self.model.lm_head def set_output_embeddings(self, new_embeddings: nn.Linear) -> None: self.model.lm_head = new_embeddings def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ipa_vectors: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, **kwargs, ) -> CausalLMOutput: B, T = input_ids.shape if ipa_vectors is None: ipa_vectors = torch.zeros( B, T, self.config.ipa_dim, dtype=torch.float32, device=input_ids.device, ) out = self.model(input_ids, ipa_vectors, attention_mask) loss = None if labels is not None: shift_logits = out.logits[:, :-1].contiguous() shift_labels = labels[:, 1:].contiguous() loss = nn.functional.cross_entropy( shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1), ignore_index=-100, ) loss = loss + self.config.ipa_lambda * out.aux_loss return CausalLMOutput(loss=loss, logits=out.logits)