#!/usr/bin/env python3
"""
Parametric HAWK for pretraining AND the full BabyLM eval pipeline.

Heads (all share ONE parametric backbone, so a CausalLM checkpoint maps 1:1 onto
every head's wte/layers/norm_f/params/film/switch tensors -- only the task head is
newly initialised):
  - HawkForCausalLM              : LM pretraining + BLiMP / MultiBLiMP (zero-shot LL)
  - HawkForSequenceClassification: SuperGLUE-style fine-tuning (pooled last non-pad token)
  - HawkForTokenClassification   : token tasks (per-token logits, ignore_index=-100)

parametric_mode:
  "film"   : InputInferredParameters router (k bits) + BCE param_sup + diffuse FiLM.
  "switch" : 3 named single-site switches, gold p teacher-forced by language, router
             kept as a detached inference predictor; wh_ex_situ->RG-LRU decay (memory horizon), v2/null_subject
             ->token-anchored expectation surrogate; force_params counterfactual.
All parametric code is gated on config.k_params>0; k_params=0 == pristine Hawk backbone.

transformers 4.x/5.x tie fix: `_tied_weights_keys` is a LIST and lm_head is HARD-TIED to
wte in __init__ (sharing the Parameter object), so v5's meta-device / separate-storage tie
resolution can't desync, and save_pretrained dedups lm_head.weight cleanly.
"""
from typing import Optional
import torch, torch.nn as nn, torch.nn.functional as F
import transformers
from transformers import PreTrainedModel, PretrainedConfig
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import (
    BaseModelOutput, CausalLMOutputWithPast, SequenceClassifierOutput, TokenClassifierOutput)

# transformers 4.x expects _tied_weights_keys as a LIST of target keys; 5.x expects a DICT
# {target: source} (5.13 hard-crashes on the list form in post_init). Adapt at import time.
_TF_MAJOR = int(transformers.__version__.split(".")[0])
_TIED_KEYS = ({"lm_head.weight": "wte.weight"} if _TF_MAJOR >= 5 else ["lm_head.weight"])

ROBERTS_TABLE = {  # film mode
    "head_dir_clausal": (1.0, 0.5, 1.0), "v2": (0.0, 1.0, 0.0),
    "null_subject": (0.0, 0.0, 1.0), "tense_morph": (1.0, 1.0, 0.0),
    "sv_agreement": (1.0, 1.0, 0.0), "filler_gap": (1.0, 1.0, 0.0),
    "head_dir_dp": (0.5, 0.5, 0.0), "def_article": (1.0, 1.0, 0.0),
    "number_morph": (1.0, 1.0, 0.0), "classifier": (0.0, 0.0, 1.0),
    "binding_domain": (0.0, 0.0, 1.0), "rc_position": (1.0, 1.0, 0.0)}
DEFAULT_SUPERVISE = ["v2", "filler_gap", "null_subject"]  # rank-3 over eng/nld/zho
SWITCH_BITS = {"v2": 0, "null_subject": 1, "wh_ex_situ": 2}
SWITCH_TABLE = {"v2": (0.0, 1.0, 0.0), "null_subject": (0.0, 0.0, 1.0), "wh_ex_situ": (1.0, 1.0, 0.0)}
SWITCH_SUPERVISE = ["v2", "null_subject", "wh_ex_situ"]


def gumbel_sigmoid(logits, tau, hard=False, training=True):
    if training:
        u = torch.rand_like(logits).clamp_(1e-6, 1 - 1e-6)
        y = torch.sigmoid((logits + torch.log(u) - torch.log1p(-u)) / tau)
    else:
        y = torch.sigmoid(logits / tau)
    if hard:
        y = (y > 0.5).float() + y - y.detach()
    return y


def diag_linear_scan(a, b):
    """Inclusive parallel scan of the diagonal affine recurrence
           h_t = a_t * h_{t-1} + b_t ,  h_0 = 0        (a, b: (B, T, W))
    via Hillis-Steele in REAL space: ceil(log2 T) vectorised passes instead of a
    T-step Python loop. Exact (matches the loop up to float rounding, and with LOWER
    accumulated error since depth is log T not T), stable because a in (0,1) keeps the
    running products <=1 (no exp(-cumsum) overflow), and static-shape so torch.compile
    traces it with no data-dependent control flow (T is a python int -> the while unrolls)."""
    T = a.shape[1]
    A, H = a, b
    d = 1
    while d < T:                       # unrolled at trace time (T static)
        one = A.new_ones(A.shape[0], d, A.shape[2])
        zero = H.new_zeros(H.shape[0], d, H.shape[2])
        A_prev = torch.cat([one,  A[:, :-d]], dim=1)   # a_{t-d}, identity 1 for t<d
        H_prev = torch.cat([zero, H[:, :-d]], dim=1)   # h_{t-d}, identity 0 for t<d
        H = A * H_prev + H
        A = A * A_prev
        d *= 2
    return H


class InputInferredParameters(nn.Module):
    def __init__(self, hidden, k=8, init_temp=1.0, logit_cap=6.0):
        super().__init__()
        self.router = nn.Sequential(nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, k))
        self.default_logits = nn.Parameter(torch.zeros(k))
        self.register_buffer("temperature", torch.tensor(float(init_temp)))
        self.logit_cap = float(logit_cap)
        nn.init.zeros_(self.router[-1].weight); nn.init.zeros_(self.router[-1].bias)
        self._last_logits = None

    def set_temperature(self, t): self.temperature.fill_(float(t))

    def forward(self, x, attn_mask=None, hard=False):
        if attn_mask is not None:
            m = attn_mask.unsqueeze(-1).to(x.dtype); pooled = (x * m).sum(1) / m.sum(1).clamp(min=1.0)
        else:
            pooled = x.mean(1)
        raw = self.router(pooled)
        logits = self.logit_cap * torch.tanh(raw / self.logit_cap) if self.logit_cap > 0 else raw
        self._last_logits = logits
        return gumbel_sigmoid(logits, self.temperature, hard, self.training)  # tensor tau: no recompile on anneal

    def ig_penalty(self):
        if self._last_logits is None:
            return self.default_logits.new_zeros(())
        return ((self._last_logits - self.default_logits.unsqueeze(0)) ** 2).mean()


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__(); self.weight = nn.Parameter(torch.ones(dim)); self.eps = eps
    def forward(self, x):
        return self.weight * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps))


class RGLRU(nn.Module):
    def __init__(self, width, c=8.0, use_parallel_scan=True):
        super().__init__()
        self.width = width; self.c = c; self.use_parallel_scan = use_parallel_scan
        self.input_gate = nn.Linear(width, width); self.recur_gate = nn.Linear(width, width)
        self.log_lambda = nn.Parameter(torch.empty(width).uniform_(2.197, 6.907))

    def forward(self, x, recur_bias=None, decay_scale=None, slot=None):
        B, T, W = x.shape
        pre_r = self.recur_gate(x)
        if recur_bias is not None:
            pre_r = pre_r + recur_bias.unsqueeze(1)
        r = torch.sigmoid(pre_r)
        i = torch.sigmoid(self.input_gate(x))
        log_a = -F.softplus(-self.log_lambda)
        log_a_t = self.c * r * log_a
        if decay_scale is not None:
            log_a_t = log_a_t * decay_scale.unsqueeze(1)
        if slot is not None:
            # wh-slot: (cmask (W,) 1 on protected channels, hold_eps, wh_t (B,T,1), reset_t (B,T,1))
            cmask, hold_eps, wh_t, reset_t = slot
            log_a_t = log_a_t * (1.0 - cmask + cmask * hold_eps)   # hold: a -> ~1 on the slot
        a_t = torch.exp(log_a_t)
        if slot is not None:
            a_t = a_t * (1.0 - cmask * reset_t)                    # boundary reset (PIC), gated upstream
        mult = torch.sqrt(torch.clamp(-torch.expm1(2.0 * log_a_t), min=1e-8))
        gated_x = mult * (i * x)
        if slot is not None:
            # slot writes: only at wh tokens, UNNORMALISED (the sqrt(1-a^2) factor is ~0 on a
            # held channel and would crush the write -- the wh_gamma dead-at-init failure mode)
            gated_x = gated_x * (1.0 - cmask) + cmask * (i * x) * wh_t
        if self.use_parallel_scan:
            return diag_linear_scan(a_t, gated_x)          # ceil(log2 T) parallel passes
        h = torch.zeros(B, W, device=x.device, dtype=x.dtype); outs = []   # exact sequential fallback
        for t in range(T):
            h = a_t[:, t] * h + gated_x[:, t]; outs.append(h)
        return torch.stack(outs, 1)


class RecurrentBlock(nn.Module):
    def __init__(self, d_model, d_rnn, conv_kernel=4, rglru_c=8.0, use_parallel_scan=True):
        super().__init__()
        self.conv_kernel = conv_kernel
        self.in_gate = nn.Linear(d_model, d_rnn); self.in_recur = nn.Linear(d_model, d_rnn)
        self.conv = nn.Conv1d(d_rnn, d_rnn, conv_kernel, groups=d_rnn, padding=conv_kernel - 1)
        self.rglru = RGLRU(d_rnn, rglru_c, use_parallel_scan); self.out = nn.Linear(d_rnn, d_model)
    def forward(self, x, recur_bias=None, decay_scale=None, slot=None):
        gate = F.gelu(self.in_gate(x))
        rec = self.in_recur(x).transpose(1, 2)
        rec = self.conv(rec)[..., : x.size(1)]
        rec = self.rglru(rec.transpose(1, 2), recur_bias, decay_scale, slot)
        return self.out(gate * rec)


class MLPBlock(nn.Module):
    def __init__(self, d_model, expansion=3):
        super().__init__(); hidden = expansion * d_model
        self.gate = nn.Linear(d_model, hidden); self.up = nn.Linear(d_model, hidden)
        self.down = nn.Linear(hidden, d_model)
    def forward(self, x): return self.down(F.gelu(self.gate(x)) * self.up(x))


class HawkLayer(nn.Module):
    def __init__(self, d_model, d_rnn, conv_kernel, mlp_expansion, eps, rglru_c=8.0, use_parallel_scan=True):
        super().__init__()
        self.norm1 = RMSNorm(d_model, eps)
        self.recur = RecurrentBlock(d_model, d_rnn, conv_kernel, rglru_c, use_parallel_scan)
        self.norm2 = RMSNorm(d_model, eps); self.mlp = MLPBlock(d_model, mlp_expansion)
    def forward(self, x, recur_bias=None, decay_scale=None, slot=None):
        x = x + self.recur(self.norm1(x), recur_bias, decay_scale, slot)
        x = x + self.mlp(self.norm2(x)); return x


class HawkConfig(PretrainedConfig):
    model_type = "hawk_rglru"
    def __init__(self, vocab_size=16384, n_layer=12, n_embd=768, rnn_width=None, conv_kernel=4,
                 mlp_expansion=3, rmsnorm_eps=1e-6, rglru_c=8.0, max_position_embeddings=1024,
                 tie_word_embeddings=True, bos_token_id=2, eos_token_id=3, pad_token_id=1,
                 classifier_dropout=0.1, use_parallel_scan=True, k_params=0, parametric_mode="film",
                 param_inject="embed",
                 boundary_token_ids=None, wh_decay_gamma=0.5,
                 wh_slot_channels=0, wh_token_ids=None, slot_hold_eps=0.02,
                 lambda_param_sup=0.0, param_reg=1e-2,
                 param_label_smooth=0.05, logit_cap=6.0, gumbel_temperature=1.0, **kwargs):
        self.vocab_size = vocab_size; self.n_layer = n_layer; self.n_embd = n_embd
        self.rnn_width = rnn_width; self.conv_kernel = conv_kernel; self.mlp_expansion = mlp_expansion
        self.rmsnorm_eps = rmsnorm_eps; self.rglru_c = rglru_c
        self.max_position_embeddings = max_position_embeddings; self.classifier_dropout = classifier_dropout
        self.use_parallel_scan = use_parallel_scan
        self.k_params = k_params; self.parametric_mode = parametric_mode; self.param_inject = param_inject
        self.boundary_token_ids = boundary_token_ids or []; self.wh_decay_gamma = wh_decay_gamma
        self.wh_slot_channels = wh_slot_channels; self.wh_token_ids = wh_token_ids or []
        self.slot_hold_eps = slot_hold_eps
        self.lambda_param_sup = lambda_param_sup; self.param_reg = param_reg
        self.param_label_smooth = param_label_smooth; self.logit_cap = logit_cap
        self.gumbel_temperature = gumbel_temperature
        if parametric_mode == "switch" and k_params > 0:
            assert k_params == 3, "switch mode uses exactly 3 named bits (set k_params=3)"
        self.auto_map = {
            "AutoConfig": "modeling_hawk_param.HawkConfig",
            # AutoModel = parametric backbone only (last_hidden_state). Required by
            # finetune_token_classification.py, which wraps AutoModel in its own
            # PooledTokenClassifier rather than using AutoModelForTokenClassification.
            "AutoModel": "modeling_hawk_param.HawkModel",
            "AutoModelForCausalLM": "modeling_hawk_param.HawkForCausalLM",
            "AutoModelForSequenceClassification": "modeling_hawk_param.HawkForSequenceClassification",
            "AutoModelForTokenClassification": "modeling_hawk_param.HawkForTokenClassification"}
        super().__init__(tie_word_embeddings=tie_word_embeddings, bos_token_id=bos_token_id,
                         eos_token_id=eos_token_id, pad_token_id=pad_token_id, **kwargs)
    @property
    def d_rnn(self): return self.rnn_width if self.rnn_width is not None else self.n_embd

    # Generic HF pipeline code reads config.hidden_size; HawkConfig calls it
    # n_embd. Alias it (read/write, so from_pretrained can set it if present).
    @property
    def hidden_size(self): return self.n_embd

    @hidden_size.setter
    def hidden_size(self, value): self.n_embd = value


# --------------------------------------------------------------------------------------- #
#  Shared parametric backbone. Every head keeps wte/layers/norm_f (and the parametric
#  modules) as TOP-LEVEL attributes with identical names -> 1:1 state_dict mapping.
# --------------------------------------------------------------------------------------- #
class HawkPreTrainedModel(PreTrainedModel):
    config_class = HawkConfig
    supports_gradient_checkpointing = False

    def _init_backbone(self, config):
        H, d_rnn = config.n_embd, config.d_rnn
        self.wte = nn.Embedding(config.vocab_size, H)
        self.layers = nn.ModuleList([HawkLayer(H, d_rnn, config.conv_kernel, config.mlp_expansion,
                                               config.rmsnorm_eps, config.rglru_c,
                                               config.use_parallel_scan)
                                     for _ in range(config.n_layer)])
        self.norm_f = RMSNorm(H, config.rmsnorm_eps)

    def _init_params(self, config):
        H, d_rnn = config.n_embd, config.d_rnn
        self.params = None; self.film_embed = None; self.film_recur = None
        self.wh_gamma = None; self.event_expect = None
        if config.k_params > 0:
            self.params = InputInferredParameters(H, config.k_params, config.gumbel_temperature, config.logit_cap)
            if config.parametric_mode == "switch":
                self.wh_gamma = nn.Parameter(torch.full((d_rnn,), float(config.wh_decay_gamma)))
                self.event_expect = nn.Parameter(torch.randn(2, 2, H) * 0.02)
            else:
                if config.param_inject in ("embed", "both"):
                    self.film_embed = nn.Linear(config.k_params, H, bias=False); nn.init.normal_(self.film_embed.weight, std=0.02)
                if config.param_inject in ("rglru", "both"):
                    self.film_recur = nn.Linear(config.k_params, d_rnn, bias=False); nn.init.normal_(self.film_recur.weight, std=0.02)
        self.register_buffer("param_targets", torch.zeros(0), persistent=False)
        self.register_buffer("param_sup_mask", torch.zeros(0), persistent=False)
        self._aux = {}; self._last_p = None

    # --- parametric API (shared) ---
    def get_input_embeddings(self): return self.wte
    def set_input_embeddings(self, new): self.wte = new
    def set_temperature(self, t):
        if getattr(self, "params", None) is not None:
            self.params.set_temperature(t); self.config.gumbel_temperature = float(t)

    def set_param_targets(self, langs, table=None, supervise=None, bit_assignment=None):
        if getattr(self, "params", None) is None: return {}
        if self.config.parametric_mode == "switch":
            table = table or SWITCH_TABLE; supervise = supervise or SWITCH_SUPERVISE
            bit_assignment = bit_assignment or SWITCH_BITS
        table = table or ROBERTS_TABLE; supervise = supervise or DEFAULT_SUPERVISE
        if bit_assignment is None: bit_assignment = {n: i for i, n in enumerate(supervise)}
        k = self.config.k_params; tgt = torch.zeros(len(langs), k); mask = torch.zeros(k)
        for name in supervise:
            b = bit_assignment[name]; mask[b] = 1.0
            for li in range(len(langs)): tgt[li, b] = float(table[name][li])
        dev = self.wte.weight.device
        self.param_targets = tgt.to(dev); self.param_sup_mask = mask.to(dev)
        return {lg: i for i, lg in enumerate(langs)}

    def aux_loss(self):
        if not self._aux: return self.wte.weight.new_zeros(())
        c = self.config
        return (c.param_reg * self._aux.get("param", 0.0) + c.lambda_param_sup * self._aux.get("param_sup", 0.0))

    def aux_components(self):
        return {k: (float(v.detach()) if torch.is_tensor(v) else float(v)) for k, v in self._aux.items()}

    def _backbone(self, input_ids, attention_mask=None, lang_ids=None, force_params=None):
        """embed -> (infer + inject typological params) -> layers -> norm_f. Sets self._aux."""
        cfg = self.config; x = self.wte(input_ids); B = x.size(0)
        self._aux = {}; recur_bias = None; decay_scale = None
        sw = (cfg.parametric_mode == "switch")
        if getattr(self, "params", None) is not None:
            if attention_mask is None and cfg.pad_token_id is not None:
                attention_mask = (input_ids != cfg.pad_token_id).long()
            p_router = self.params(x, attention_mask)
            if sw:
                if lang_ids is not None and self.param_targets.numel() > 0:
                    p = self.param_targets[lang_ids]
                else:
                    p = (p_router > 0.5).to(x.dtype).detach()
            else:
                p = p_router
            if force_params is not None:
                p = p.clone()
                if isinstance(force_params, dict):
                    for kb, v in force_params.items():
                        p[:, SWITCH_BITS[kb] if isinstance(kb, str) else int(kb)] = float(v)
                else:
                    p = torch.as_tensor(force_params, dtype=p.dtype, device=p.device).expand_as(p).clone()
            self._last_p = p          # grad-carrying: the wh-slot escape gate trains through it
            self._last_p_detached = p.detach()
            if sw:
                p_wh = p[:, 2:3]
                decay_scale = torch.exp(self.wh_gamma.unsqueeze(0) * (1.0 - 2.0 * p_wh))
                E = self.event_expect
                # bit 0 (v2): expectation about clause POSITION 2 (finite verb second).
                # bit 1 (null_subject): expectation about clause POSITION 1 (overt argument
                # required vs droppable). The boundary token's step predicts position 1 ->
                # null_subject injects there (bmask); the position-1 token's step predicts
                # position 2 -> v2 injects there (amask). Injections also seed the RG-LRU
                # state, so their influence persists into the clause with the decay horizon.
                e_v2 = p[:, 0:1] * E[0, 1] + (1 - p[:, 0:1]) * E[0, 0]
                e_ns = p[:, 1:2] * E[1, 1] + (1 - p[:, 1:2]) * E[1, 0]
                if cfg.boundary_token_ids:
                    bids = torch.as_tensor(cfg.boundary_token_ids, device=input_ids.device)
                    bmask = torch.isin(input_ids, bids).to(x.dtype)      # predicts clause pos 1
                    amask = F.pad(bmask, (1, 0))[:, :-1]                 # predicts clause pos 2
                    x = x + bmask.unsqueeze(-1) * e_ns.unsqueeze(1) + amask.unsqueeze(-1) * e_v2.unsqueeze(1)
                else:
                    x = x + e_v2.unsqueeze(1) + e_ns.unsqueeze(1)
            else:
                if self.film_embed is not None: x = x + self.film_embed(p).unsqueeze(1)
                if self.film_recur is not None: recur_bias = self.film_recur(p)
            self._aux["param"] = self.params.ig_penalty()
            if (cfg.lambda_param_sup > 0 and lang_ids is not None and self.param_targets.numel() > 0):
                s = float(cfg.param_label_smooth)
                tgt = self.param_targets[lang_ids] * (1 - 2 * s) + s
                bce = F.binary_cross_entropy_with_logits(self.params._last_logits, tgt, reduction="none")
                msk = self.param_sup_mask.unsqueeze(0)
                self._aux["param_sup"] = (bce * msk).sum() / (msk.sum() * B).clamp(min=1.0)
        slot = None
        if cfg.wh_slot_channels > 0 and cfg.wh_token_ids:
            K = min(cfg.wh_slot_channels, cfg.d_rnn)
            cmask = x.new_zeros(cfg.d_rnn); cmask[-K:] = 1.0        # slot = last K RG-LRU channels
            wids = torch.as_tensor(cfg.wh_token_ids, device=input_ids.device)
            wh_t = torch.isin(input_ids, wids).to(x.dtype).unsqueeze(-1)        # (B,T,1) writes
            reset_t = x.new_zeros(input_ids.shape + (1,))
            if cfg.boundary_token_ids:
                bids = torch.as_tensor(cfg.boundary_token_ids, device=input_ids.device)
                bnd = torch.isin(input_ids, bids).to(x.dtype).unsqueeze(-1)     # (B,T,1)
                # PIC with parameterised escape: filler_gap bit (film bit 1) holds the slot
                # across the boundary (eng/nld); its absence resets it there (zho).
                if self._last_p is not None and self._last_p.size(1) > 1:
                    p_fg = self._last_p[:, 1:2].to(x.dtype)
                else:
                    p_fg = x.new_zeros(x.size(0), 1)
                reset_t = bnd * (1.0 - p_fg)[:, None, :]
            slot = (cmask, float(cfg.slot_hold_eps), wh_t, reset_t)
        for layer in self.layers:
            x = layer(x, recur_bias, decay_scale, slot)
        return self.norm_f(x)


class HawkModel(HawkPreTrainedModel):
    """Parametric backbone only: returns `last_hidden_state`, no LM/task head.

    finetune_token_classification.py builds its own head on top of
    `AutoModel.from_pretrained(...)` and reads `outputs.last_hidden_state`;
    HawkForCausalLM returns logits, so it cannot serve that role. Shares
    _init_backbone/_init_params/_backbone with every other head, so a CausalLM
    checkpoint maps 1:1 and nothing is newly initialised. The typological
    switchboard still runs (it lives in _backbone), so POS fine-tuning sees the
    same parametric representation the LM was trained with.
    """

    def __init__(self, config: HawkConfig):
        super().__init__(config)
        self._init_backbone(config); self._init_params(config)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, lang_ids=None,
                force_params=None, **kwargs):
        x = self._backbone(input_ids, attention_mask, lang_ids, force_params)
        return BaseModelOutput(last_hidden_state=x)
class HawkForCausalLM(HawkPreTrainedModel, GenerationMixin):
    _tied_weights_keys = _TIED_KEYS                # dict on tf>=5, list on 4.x

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        """transformers 5.x can leave lm_head on the meta device (tie declared but not
        executed: seen on 5.3.0 with only wte.weight in the file, and whenever BOTH weights
        are present the loader refuses to tie). Weights are hard-tied at training time, so
        unconditionally re-aliasing lm_head to wte after load is always correct."""
        model = super().from_pretrained(*args, **kwargs)
        lm = getattr(model, "lm_head", None)
        if lm is not None:
            if lm.weight.device.type == "meta" or torch.equal(lm.weight, model.wte.weight):
                model.lm_head.weight = model.wte.weight
            elif lm.weight is not model.wte.weight:
                raise ValueError(
                    "lm_head.weight differs from wte.weight but the checkpoint declares "
                    "tie_word_embeddings=True — this checkpoint was trained UNTIED; "
                    "re-aliasing would discard a trained output head.")
        return model

    def __init__(self, config: HawkConfig):
        super().__init__(config)
        self._init_backbone(config); self._init_params(config)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight       # HARD tie (same Parameter object)
        self.post_init()

    def get_output_embeddings(self): return self.lm_head
    def set_output_embeddings(self, new): self.lm_head = new

    def forward(self, input_ids, attention_mask=None, labels=None, lang_ids=None,
                force_params=None, **kwargs):
        x = self._backbone(input_ids, attention_mask, lang_ids, force_params)
        logits = self.lm_head(x)
        loss = None
        if labels is not None:
            sl = logits[:, :-1, :].contiguous(); st = labels[:, 1:].contiguous()
            loss = F.cross_entropy(sl.view(-1, sl.size(-1)), st.view(-1), ignore_index=-100)
        return CausalLMOutputWithPast(loss=loss, logits=logits)


class HawkForSequenceClassification(HawkPreTrainedModel):
    """SuperGLUE-style fine-tuning. Pools the last non-pad hidden state (right padding).
    Backbone (incl. parametric modules) maps 1:1 from a CausalLM checkpoint; only `score`
    is newly initialised. Task loss only -- the pretraining aux terms are not folded in."""
    def __init__(self, config: HawkConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        self._init_backbone(config); self._init_params(config)
        self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, labels=None, lang_ids=None, **kwargs):
        x = self._backbone(input_ids, attention_mask, lang_ids)
        logits = self.score(x)                                   # (B,T,num_labels)
        B, T = input_ids.shape[:2]
        if attention_mask is not None:
            last = attention_mask.long().sum(-1) - 1
        elif self.config.pad_token_id is not None:
            last = (input_ids != self.config.pad_token_id).int().sum(-1) - 1
        else:
            last = torch.full((B,), T - 1, device=input_ids.device)
        last = last.clamp(min=0)
        pooled = logits[torch.arange(B, device=input_ids.device), last]   # (B,num_labels)
        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and labels.dtype in (torch.long, torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"
            if self.config.problem_type == "regression":
                lf = nn.MSELoss()
                loss = lf(pooled.squeeze(), labels.squeeze()) if self.num_labels == 1 else lf(pooled, labels)
            elif self.config.problem_type == "single_label_classification":
                loss = nn.CrossEntropyLoss()(pooled.view(-1, self.num_labels), labels.view(-1))
            else:
                loss = nn.BCEWithLogitsLoss()(pooled, labels.float())
        return SequenceClassifierOutput(loss=loss, logits=pooled)


class HawkForTokenClassification(HawkPreTrainedModel):
    """Per-token classification (POS / morphology). Standard ignore_index=-100 loss."""
    def __init__(self, config: HawkConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        self._init_backbone(config); self._init_params(config)
        drop = config.classifier_dropout if config.classifier_dropout is not None else 0.1
        self.dropout = nn.Dropout(drop)
        self.classifier = nn.Linear(config.n_embd, self.num_labels)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, labels=None, lang_ids=None, **kwargs):
        x = self._backbone(input_ids, attention_mask, lang_ids)
        logits = self.classifier(self.dropout(x))               # (B,T,num_labels)
        loss = None
        if labels is not None:
            loss = nn.CrossEntropyLoss()(logits.view(-1, self.num_labels), labels.view(-1))
        return TokenClassifierOutput(loss=loss, logits=logits)