hf_wrapper.py

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