AutoModel: safetensors + config + modeling + card

README.md

@@ -0,0 +1,119 @@
+---
+license: mit
+language: [en]
+library_name: transformers
+pipeline_tag: feature-extraction
+tags: [sentence-similarity, feature-extraction, consensus-distillation, geometric-deep-learning, amoe]
+datasets: [AbstractPhil/conceptual-captions-12m-webdataset-berts]
+base_model: [google-bert/bert-base-uncased, answerdotai/ModernBERT-base, FacebookAI/roberta-base, albert/albert-base-v2, distilbert/distilbert-base-uncased]
+---
+
+# captionbert-8192-v2
+
+A **58.3M** standalone sentence encoder distilled from the geometric **consensus**
+of five BERT-family teachers. No expert models at inference: tokenizer + this
+model, 768-d L2-normalized output.
+
+12 layers, 512-d, 8 heads, FFN 2048, 8192 position capacity. **0.53x bert-base.**
+
+```python
+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")
+
+emb = model.encode(["a cat on a windowsill", "a feline by the window"])   # (2, 768)
+(emb[0] @ emb[1]).item()
+```
+
+## Results
+
+Measured in one harness; every model mean-pooled and L2-normalized, no task
+tuning. `erank` is the participation ratio of the embedding spectrum -- how many
+of the 768 directions are actually used.
+
+| model | params | STS-B rho | SICK-R rho | self_cos | erank |
+|---|---|---|---|---|---|
+| bert-base | 109.5M | .4729 | .5865 | +.580 | 32.0 |
+| ModernBERT-base | 149.0M | .4215 | .5479 | +.948 | -- |
+| roberta-base | 124.6M | .5436 | .6296 | +.976 | -- |
+| albert-base-v2 | 11.7M | .4784 | .5364 | +.905 | -- |
+| distilbert | 66.4M | .5717 | .6424 | +.840 | -- |
+| **captionbert-8192-v2** | **58.3M** | **.5747** | **.6526** | **+.129** | 33.4 |
+| all-MiniLM-L6-v2 (ref) | 22.7M | .8203 | .7758 | +.023 | 94.3 |
+
+**It edges every teacher it was distilled from**, at 13% of their combined
+parameters, having never seen a similarity label. It does **not** reach
+`all-MiniLM-L6-v2`, which was contrastively trained on 1B+ curated pairs --
+a different comparison class.
+
+**Isotropy is the mechanism.** Mean-pooled BERT-family embeddings sit in a narrow
+cone (self_cos .58-.98); this model reads **+.129**, and cosine discriminates far
+better in a space that is not collapsed.
+
+## With an AMOE anchor (`amoe/`)
+
+The trunk is frozen; adapters are 1.6M params each. Anchors ship in this repo
+under `amoe/` and toggle **bit-exact** -- all anchors disabled reproduces the
+bare trunk exactly, so one artifact serves both.
+
+| config | STS-B rho | SICK-R rho |
+|---|---|---|
+| bare trunk | .5747 | .6526 |
+| + `equiv` anchor (all-nli) | .7254 | **.7550** |
+| + `simplify` anchor (wiki/altlex/compression) | .7400 | .7075 |
+| + **2-anchor dispatch (MOE)** | **.7524** | .7380 |
+
+The two anchors are complementary along the task axis, and the router separated
+them: mean `|w/z|` moved from .310/.380 (blend) to .645/.223 (specialize) over
+800 keys-only steps. See [amoe-lora](https://github.com/AbstractEyes/amoe-lora).
+
+## How it was built
+
+1. Five teachers embedded 33M CC12M llava-next captions (mean-pooled, 768-d).
+2. One global **whitened Procrustes** map per teacher into `bert-base`'s frame,
+   fit on a stratified random sample and **reported out-of-sample**.
+3. Consensus = normalized centroid of the aligned teachers, per chunk.
+4. Student trained from scratch: InfoNCE(T=0.07) + per-sample MSE against the
+   consensus. Pure Adam, no weight decay. 26.9M rows, 52,548 steps at batch
+   2048, ~5.4 h on one RTX 6000 Pro.
+
+## Known limits -- read before using
+
+- **Consensus rank is ~28.7 of 768.** Five BERT-family teachers only agree on
+  about 29 directions. The student uses ~103 in domain but falls back to ~33 on
+  out-of-domain text: **the structure it builds on captions does not transfer.**
+  This is the model's ceiling and it is a property of the consensus, not the
+  student.
+- **Alignment quality varies by teacher.** Out-of-sample cosine to the bert
+  frame: distil .625, roberta .372, albert .331, modern .327. The ordering
+  tracks architectural distance from bert-base.
+- **10 of 66 source chunks lacked ModernBERT**, so 54 chunks (~27M rows) were
+  used. No 4-expert fallback: that would change the target definition mid-dataset.
+- Trained on image captions. Expect caption-like text to be its strongest domain.
+- Single seed. No variance estimate.
+
+## Output convention (differs from v1)
+
+| field | shape | |
+|---|---|---|
+| `last_hidden_state` | (B, L, 512) | token states |
+| `pooler_output` | (B, 768) | **the embedding**, L2-normalized |
+| `embedding` | (B, 768) | alias |
+
+`geolip-captionbert-8192` (v1) returned the pooled embedding as
+`last_hidden_state`. If porting v1 code, use `pooler_output`. v1 also shipped an
+`AlignmentBank`; v2 does not -- measured on v1, its expert-consistency features
+varied 0.2% across samples because a rotation round-trip carries no data.
+
+## Citation
+
+```bibtex
+@misc{abstractphil2026captionbertv2,
+  title  = {captionbert-8192-v2: consensus distillation at CC12M scale},
+  author = {AbstractPhil},
+  year   = {2026},
+  url    = {https://huggingface.co/AbstractPhil/captionbert-8192-v2}
+}
+```
+
+MIT.

config.json

@@ -0,0 +1,22 @@
+{
+  "architectures": [
+    "CaptionBertV2Model"
+  ],
+  "model_type": "captionbert_v2",
+  "auto_map": {
+    "AutoConfig": "modeling_captionbert.CaptionBertV2Config",
+    "AutoModel": "modeling_captionbert.CaptionBertV2Model"
+  },
+  "vocab_size": 30522,
+  "hidden_size": 512,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 8,
+  "intermediate_size": 2048,
+  "output_dim": 768,
+  "max_position_embeddings": 8192,
+  "hidden_dropout_prob": 0.1,
+  "pad_token_id": 0,
+  "pooling": "mean",
+  "torch_dtype": "float32",
+  "tokenizer_class": "BertTokenizerFast"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:303bfdbe4c2e0a920124849a5b72226afad6091172e05c74332d6a31c458f5f3
+size 233251232

modeling_captionbert.py

@@ -1,210 +1,210 @@
-# ============================================================================
-# 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"])
-#
-# ---------------------------------------------------------------------------
-# 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.
-# ============================================================================
-
-from dataclasses import dataclass
-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)
-
-
-CaptionBertV2Config.register_for_auto_class()
+# ============================================================================
+# 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"])
+#
+# ---------------------------------------------------------------------------
+# 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.
+# ============================================================================
+
+from dataclasses import dataclass
+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)
+
+
+CaptionBertV2Config.register_for_auto_class()
 CaptionBertV2Model.register_for_auto_class("AutoModel")