modeling_captionbert.py

# ============================================================================
# modeling_captionbert.py -- AbstractPhil/captionbert-8192-v2
#
#   from transformers import AutoModel, AutoTokenizer
#   model = AutoModel.from_pretrained("AbstractPhil/captionbert-8192-v2",
#                                     trust_remote_code=True)
#   tok   = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
#   out   = model(**tok(["a cat on a windowsill"], return_tensors="pt"))
#   emb   = out.pooler_output          # (B, 768) L2-normalized
#
# OR just:  emb = model.encode(["a cat on a windowsill"])
#
# WITH AMOE ARMS (needs amoe-lora; imported lazily so the base model loads
# without it):
#   model.attach_amoe()                # shipped 3-arm collective
#   emb = model.encode([...])          # adapted
#   with model.amoe_off():             # the unsupervised baseline
#       base = model.encode([...])
#   model.set_amoe(["equiv"])          # one arm, DAMPED inside the dispatch
#   model.detach_amoe()                # bit-exact restore, asserted
#
# ---------------------------------------------------------------------------
# BREAKING CHANGE FROM v1 -- READ THIS IF YOU USED geolip-captionbert-8192
#   v1 returned the POOLED 768-d embedding as `last_hidden_state`. That is not
#   the transformers convention and it silently breaks anything expecting token
#   states. v2 follows the convention:
#       last_hidden_state : (B, L, 512)  token states
#       pooler_output     : (B, 768)     L2-normalized embedding  <-- the product
#       embedding         : (B, 768)     alias for pooler_output
#   If you are porting v1 code, `last_hidden_state` -> `pooler_output`.
#
#   v1 also shipped an AlignmentBank. v2 does NOT. Measured on v1: the bank's
#   expert-consistency block varied 0.2% across samples and took 0.23% of its
#   projection energy while anchor distances took 98.70% -- because
#   `back = x @ R.T @ R` is a rotation round-trip and carries no data. Content
#   extensions belong in an AMOE anchor, which this repo ships separately.
# ============================================================================

import os
from dataclasses import dataclass
from types import SimpleNamespace
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPooling


class CaptionBertV2Config(PretrainedConfig):
    model_type = "captionbert_v2"

    def __init__(
        self,
        vocab_size: int = 30522,
        hidden_size: int = 512,          # d_model
        num_hidden_layers: int = 12,
        num_attention_heads: int = 8,
        intermediate_size: int = 2048,
        output_dim: int = 768,           # consensus space
        max_position_embeddings: int = 8192,
        hidden_dropout_prob: float = 0.1,
        pad_token_id: int = 0,
        pooling: str = "mean",           # "mean" | "cls"
        **kwargs,
    ):
        super().__init__(pad_token_id=pad_token_id, **kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.output_dim = output_dim
        self.max_position_embeddings = max_position_embeddings
        self.hidden_dropout_prob = hidden_dropout_prob
        self.pooling = pooling


class CaptionBertV2Model(PreTrainedModel):
    """
    Standalone caption/sentence encoder distilled from the geometric consensus
    of five BERT-family teachers. No expert models at inference.

    Parameter names are deliberately NOT namespaced under a submodule so that
    the training checkpoint loads unchanged: token_emb, pos_emb, emb_norm,
    encoder.layers.*, output_proj.*.
    """

    config_class = CaptionBertV2Config
    base_model_prefix = "captionbert_v2"
    supports_gradient_checkpointing = True

    def __init__(self, config: CaptionBertV2Config):
        super().__init__(config)
        d = config.hidden_size
        self.token_emb = nn.Embedding(config.vocab_size, d,
                                      padding_idx=config.pad_token_id)
        self.pos_emb = nn.Embedding(config.max_position_embeddings, d)
        self.emb_norm = nn.LayerNorm(d)
        self.emb_drop = nn.Dropout(config.hidden_dropout_prob)
        layer = nn.TransformerEncoderLayer(
            d_model=d,
            nhead=config.num_attention_heads,
            dim_feedforward=config.intermediate_size,
            dropout=config.hidden_dropout_prob,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(
            layer, num_layers=config.num_hidden_layers, enable_nested_tensor=False)
        self.output_proj = nn.Sequential(
            nn.Linear(d, d), nn.GELU(), nn.LayerNorm(d), nn.Linear(d, config.output_dim))
        self.post_init()

    # -- HF plumbing --
    def get_input_embeddings(self):
        return self.token_emb

    def set_input_embeddings(self, value):
        self.token_emb = value

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        return_dict = return_dict if return_dict is not None else True
        L = input_ids.shape[1]
        pos = torch.arange(L, device=input_ids.device).unsqueeze(0)
        x = self.emb_drop(self.emb_norm(self.token_emb(input_ids) + self.pos_emb(pos)))

        kpm = (~attention_mask.bool()) if attention_mask is not None \
            else (input_ids == self.config.pad_token_id)

        hidden = [x] if output_hidden_states else None
        # Iterate the layers directly rather than calling self.encoder(...):
        # nn.TransformerEncoder's fast path inspects layer types, and an AMOE
        # anchor wraps each layer in a BlockWithAdapter that is not a
        # TransformerEncoderLayer. This keeps attach() a drop-in.
        for mod in self.encoder.layers:
            x = mod(x, src_key_padding_mask=kpm)
            if output_hidden_states:
                hidden.append(x)
        if self.encoder.norm is not None:
            x = self.encoder.norm(x)

        if self.config.pooling == "cls":
            pooled = x[:, 0]
        else:
            m = (attention_mask.unsqueeze(-1).to(x.dtype) if attention_mask is not None
                 else (~kpm).unsqueeze(-1).to(x.dtype))
            pooled = (x * m).sum(1) / m.sum(1).clamp(min=1)
        embedding = F.normalize(self.output_proj(pooled), dim=-1)

        if not return_dict:
            return (x, embedding) + ((tuple(hidden),) if output_hidden_states else ())
        out = BaseModelOutputWithPooling(
            last_hidden_state=x,                       # (B, L, 512) token states
            pooler_output=embedding,                   # (B, 768) THE PRODUCT
            hidden_states=tuple(hidden) if output_hidden_states else None,
        )
        out.embedding = embedding                      # explicit alias
        return out

    @torch.no_grad()
    def encode(self, texts, tokenizer=None, batch_size: int = 128,
               max_length: int = 256, device=None) -> torch.Tensor:
        """Raw text -> (N, 768) L2-normalized embeddings."""
        if isinstance(texts, str):
            texts = [texts]
        if tokenizer is None:
            from transformers import AutoTokenizer
            tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
        device = device or next(self.parameters()).device
        was_training = self.training
        self.eval()
        out = []
        for i in range(0, len(texts), batch_size):
            t = tokenizer(list(texts[i:i + batch_size]), max_length=max_length,
                          padding=True, truncation=True, return_tensors="pt").to(device)
            out.append(self(**t).pooler_output.float().cpu())
        if was_training:
            self.train()
        return torch.cat(out)


# ---------------------------------------------------------------------------
# AMOE binding -- lets amoe-lora attach anchors to this trunk unmodified.
#
#   import amoe
#   from modeling_captionbert import CaptionBertV2Binding
#   h = amoe.attach(model, "amoe/moe/equiv.anchor.pt",
#                   binding=CaptionBertV2Binding(d=model.config.hidden_size))
#
# amoe's PathBinding would find encoder.layers by dotted path but then read
# model.config.hidden_size -- which works here because this IS a
# PretrainedConfig. The explicit binding is kept for plain-nn.Module use.
# ---------------------------------------------------------------------------

@dataclass
class CaptionBertV2Binding:
    d: int = 512
    name: str = "captionbert_v2"

    def layers(self, model):
        return model.encoder.layers

    def set_layers(self, model, new):
        model.encoder.layers = nn.ModuleList(new)

    def hidden_size(self, model) -> int:
        return int(self.d)


# ---------------------------------------------------------------------------
# AMOE ARMS -- attach / toggle / detach, from the same AutoModel object.
#
#   model = AutoModel.from_pretrained(REPO, trust_remote_code=True)
#   model.attach_amoe()                       # shipped 3-arm collective
#   emb = model.encode(["a cat on a windowsill"])
#   with model.amoe_off():                    # the unsupervised baseline
#       base = model.encode(["a cat on a windowsill"])
#   model.set_amoe(["equiv"])                 # one arm inside the dispatch
#   model.detach_amoe()                       # bit-exact restore, asserted
#
# `amoe` is imported LAZILY: this file is executed by every
# AutoModel.from_pretrained(trust_remote_code=True), and the base model must
# load on a machine that has never heard of amoe-lora.
#
# Masking never renormalizes (the damping law), so a single arm inside the
# dispatch is DAMPED and reads lower than that arm trained alone. That is the
# mechanism, not a bug -- see the model card.
# ---------------------------------------------------------------------------

# AMOE_REPO defaults to THE REPO THIS MODEL WAS LOADED FROM, not a hardcoded
# name. Anchors are TRUNK-BOUND: measured 2026-08-02, the captionbert-8192-v2
# arms lose 31% of their gain on the -B trunk (.7287 -> .6863), and re-aligning
# the routing keys alone recovers only 29% of that. Retrained natively they
# reach .7294. So a hardcoded arms repo silently hands a trunk the wrong arms.
AMOE_REPO = None            # None -> config._name_or_path
AMOE_LIB = "amoe/arms"
# dispatch checkpoints are searched in order; the first that exists wins
AMOE_DISPATCH_CANDIDATES = (
    "amoe/b-collective/captionbert-b-arms-native.dispatch.pt",
    "amoe/collective/captionbert-v2-collective.dispatch.pt",
)
AMOE_DEFAULT_DISPATCH = None   # None -> first candidate present in the repo
AMOE_FALLBACKS = lambda n: [f"amoe/b-collective/{n}.anchor.pt",
                            f"amoe/collective/{n}.anchor.pt",
                            f"amoe/moe/{n}.anchor.pt"]
_AMOE_HINT = ("amoe-lora is required for arms: "
              "pip install git+https://github.com/AbstractEyes/amoe-lora")


def _amoe():
    try:
        from amoe.core.adapter import AdapterSpec, RelayPatchwork
        from amoe.core.dispatch import AnchorDispatch, BlockWithDispatch
        from amoe.io.checkpoint import load_anchor, load_dispatch
    except ImportError as e:
        raise ImportError(_AMOE_HINT) from e
    return SimpleNamespace(AdapterSpec=AdapterSpec, RelayPatchwork=RelayPatchwork,
                           AnchorDispatch=AnchorDispatch,
                           BlockWithDispatch=BlockWithDispatch,
                           load_anchor=load_anchor, load_dispatch=load_dispatch)


class _ArmsOff:
    def __init__(self, model):
        self.m = model

    def __enter__(self):
        self._prev = [list(d.enabled) for d in self.m._amoe_dispatches]
        for d in self.m._amoe_dispatches:
            d.enabled = [False] * len(d.enabled)
        return self.m

    def __exit__(self, *exc):
        for d, p in zip(self.m._amoe_dispatches, self._prev):
            d.enabled = list(p)


def _attach_amoe(self, arms=None, dispatch=None, repo=None, tau=None, verify=True):
    """
    Attach AMOE arms under a trained dispatch.

    arms      list of arm names. None -> the dispatch checkpoint's own roster,
              in ITS order (the routing keys are per-arm and positional, so the
              order is not cosmetic).
    dispatch  path in `repo`, a local file, or None for an untrained dispatch.
    verify    assert that all-arms-disabled reproduces the bare trunk.
    """
    from huggingface_hub import hf_hub_download, HfApi
    A = _amoe()
    if getattr(self, "_amoe_original_layers", None) is not None:
        raise RuntimeError("arms already attached; call detach_amoe() first")

    # resolve the repo: this model's own, unless told otherwise
    if repo is None:
        repo = AMOE_REPO or getattr(self.config, "_name_or_path", None)
    if not repo:
        raise ValueError("cannot determine the arms repo; pass repo=...")
    # resolve the dispatch: first candidate actually present in THAT repo
    if dispatch is None:
        dispatch = AMOE_DEFAULT_DISPATCH
    if dispatch is None:
        try:
            present = set(HfApi().list_repo_files(repo))
        except Exception:
            present = set()
        dispatch = next((c for c in AMOE_DISPATCH_CANDIDATES if c in present), None)
        if dispatch is None:
            local = os.path.isdir(str(repo))
            hint = (f"'{repo}' is a LOCAL DIRECTORY, so there is no hub repo to "
                    f"search. Pass the repo explicitly:\n"
                    f"    model.attach_amoe(repo='AbstractPhil/captionbert-8192-v2-B')"
                    if local else
                    f"Pass dispatch=<path in {repo}>, or dispatch=False to attach "
                    f"an UNTRAINED dispatch.")
            raise FileNotFoundError(
                f"no dispatch found in {repo}. Tried "
                f"{list(AMOE_DISPATCH_CANDIDATES)}.\n         {hint}")
        print(f"[amoe] dispatch: {repo}/{dispatch}")
    if dispatch is False:
        dispatch = None

    dck = None
    if dispatch is not None:
        p = dispatch if os.path.exists(dispatch) else hf_hub_download(repo, dispatch)
        dck = A.load_dispatch(p)
        if arms is None:
            arms = list(dck.meta.get("anchors", []))
        if tau is None:
            tau = float(dck.meta.get("tau", 0.1))
    if not arms:
        raise ValueError("no arms given and the dispatch names none")
    tau = 0.1 if tau is None else tau

    cks, resolved = [], []
    for name in arms:
        if os.path.exists(str(name)):
            path = str(name)
        else:
            # amoe/arms/ is the canonical library, but it is published by a
            # consolidation step that may not have run yet. Fall back to the
            # campaign folders the anchors were originally trained into, and
            # say which one answered.
            path = None
            for cand in ([f"{AMOE_LIB}/{name}.anchor.pt"] + AMOE_FALLBACKS(name)):
                try:
                    path = hf_hub_download(repo, cand)
                    resolved.append(cand)
                    break
                except Exception:
                    continue
            if path is None:
                raise FileNotFoundError(
                    f"arm '{name}' not found in {repo}: tried {AMOE_LIB}/ and "
                    f"{AMOE_FALLBACKS(name)}. Pass an explicit path, or publish "
                    f"the arm library.")
        cks.append(A.load_anchor(path))
    if resolved and not all(r.startswith(AMOE_LIB) for r in resolved):
        print(f"[amoe] resolved outside {AMOE_LIB}: "
              f"{[r for r in resolved if not r.startswith(AMOE_LIB)]}")

    d = self.config.hidden_size
    layers = list(self.encoder.layers)
    self._amoe_original_layers = layers
    dev = next(self.parameters()).device
    new, disps = [], []
    for i, layer in enumerate(layers):
        stack = nn.ModuleList()
        for ck in cks:
            a = A.RelayPatchwork(d, A.AdapterSpec())
            a.load_state_dict({k[len(f"{i}."):]: v for k, v in ck.adapters.items()
                               if k.startswith(f"{i}.")})
            for q in a.parameters():
                q.requires_grad_(False)
            stack.append(a)
        dp = A.AnchorDispatch(stack.to(dev), d,
                              emb=int(dck.meta.get("emb", 64)) if dck else 64,
                              tau=tau).to(dev)
        if dck is not None:
            with torch.no_grad():
                dp.dispatch.copy_(dck.dispatch[i]["dispatch"].to(dev))
                dp.key_proj.copy_(dck.dispatch[i]["key_proj"].to(dev))
        for q in dp.parameters():
            q.requires_grad_(False)
        disps.append(dp)
        new.append(A.BlockWithDispatch(layer, dp))
    self.encoder.layers = nn.ModuleList(new)
    self._amoe_dispatches = disps
    self._amoe_names = list(arms)

    if verify:
        ids = (torch.arange(8, device=dev).unsqueeze(0) % 7 + 1)
        am = torch.ones_like(ids)
        was = self.training
        self.eval()
        wrapped_layers = self.encoder.layers
        try:
            with torch.no_grad():
                with self.amoe_off():
                    off = self(input_ids=ids,
                               attention_mask=am).pooler_output.clone()
                # temporarily restore the bare stack for the reference forward;
                # try/finally so a raise here cannot strand the model holding
                # unwrapped layers while _amoe_dispatches still points at arms.
                self.encoder.layers = nn.ModuleList(layers)
                pure = self(input_ids=ids,
                            attention_mask=am).pooler_output.clone()
        finally:
            self.encoder.layers = wrapped_layers
            if was:
                self.train()
        gap = (off - pure).abs().max().item()
        if gap > 1e-6:
            raise RuntimeError(
                f"TOGGLE LAW VIOLATED: arms disabled shift the trunk by {gap:.3e}. "
                "Do not trust any on/off comparison from this attachment.")
    return self


def _set_amoe(self, arms):
    """Enable a subset by name (or a boolean mask in attach order)."""
    if getattr(self, "_amoe_dispatches", None) is None:
        raise RuntimeError("no arms attached")
    if arms is None:
        mask = [True] * len(self._amoe_names)
    elif all(isinstance(x, bool) for x in arms):
        mask = list(arms)
    else:
        unknown = [a for a in arms if a not in self._amoe_names]
        if unknown:
            raise ValueError(f"unknown arms {unknown}; have {self._amoe_names}")
        mask = [n in arms for n in self._amoe_names]
    for d in self._amoe_dispatches:
        d.enabled = list(mask)
    return self


def _amoe_off(self):
    if getattr(self, "_amoe_dispatches", None) is None:
        raise RuntimeError("no arms attached")
    return _ArmsOff(self)


def _detach_amoe(self, verify=True):
    """Restore the bare trunk. Bit-exact by construction; asserted by default."""
    orig = getattr(self, "_amoe_original_layers", None)
    if orig is None:
        return self
    dev = next(self.parameters()).device
    if verify:
        was = self.training
        self.eval()
        ids = (torch.arange(8, device=dev).unsqueeze(0) % 7 + 1)
        am = torch.ones_like(ids)
        with torch.no_grad():
            with self.amoe_off():
                off = self(input_ids=ids, attention_mask=am).pooler_output.clone()
        self.encoder.layers = nn.ModuleList(orig)
        with torch.no_grad():
            pure = self(input_ids=ids, attention_mask=am).pooler_output.clone()
        if was:
            self.train()
        gap = (off - pure).abs().max().item()
        if gap > 1e-6:
            raise RuntimeError(f"detach is not bit-exact: {gap:.3e}")
    else:
        self.encoder.layers = nn.ModuleList(orig)
    self._amoe_original_layers = None
    self._amoe_dispatches = None
    self._amoe_names = None
    return self


def _amoe_arms(self):
    """Names of the attached arms, in dispatch order. Empty if none."""
    return list(getattr(self, "_amoe_names", None) or [])


CaptionBertV2Model.attach_amoe = _attach_amoe
CaptionBertV2Model.set_amoe = _set_amoe
CaptionBertV2Model.amoe_off = _amoe_off
CaptionBertV2Model.detach_amoe = _detach_amoe
CaptionBertV2Model.amoe_arms = property(_amoe_arms)

CaptionBertV2Config.register_for_auto_class()
CaptionBertV2Model.register_for_auto_class("AutoModel")