MimeLens-001 release: mimelens-001-tiny-bpe-64k-s2 (paper-applied @ git 47c142d)

README.md

@@ -0,0 +1,127 @@
+---
+license: mit
+library_name: transformers
+tags:
+  - file-type-detection
+  - mime-classification
+  - binary-content
+  - position-agnostic
+  - libmagic
+  - bpe
+  - mimelens
+language: en
+base_model: mjbommar/binary-tokenizer-001-64k
+pipeline_tag: feature-extraction
+model-index:
+  - name: mimelens-001-tiny-bpe-64k-s2
+    results:
+      - task:
+          type: feature-extraction
+          name: MIME-125 classification (libmagic 125-class taxonomy)
+        dataset:
+          name: magic-bpe magic-frags (4 KB head of 64 KB random chunks, n=4,096)
+          type: mjbommar/magic-bpe-stratified
+        metrics:
+          - name: top-1 accuracy
+            type: accuracy
+            value: 0.7324
+          - name: macro-F1
+            type: f1
+            value: 0.6086
+          - name: kNN R@1
+            type: recall@1
+            value: 0.6754
+        source:
+          name: "MimeLens paper (Bommarito 2026), Appendix A"
+          url: https://github.com/mjbommar/binary-embedding-paper
+---
+
+# MimeLens-001 / tiny / bpe-64k / s2
+
+**One cell from the [MimeLens-001](https://huggingface.co/mjbommar/mimelens-001) family** — `3.15` M backbone params, `bpe-64k` input pipeline, seed `2`. Pretrained MLM-only on 33 GB of position-arbitrary binary content for fine-grained file-content-type classification under [libmagic](https://github.com/file/file)'s 125-class MIME taxonomy.
+
+A single 4 KB byte buffer in (of which the first 1,022 body tokens are consumed), one of libmagic's 125 MIME labels out, regardless of where in a source file the buffer came from.
+
+For the family overview, decision tree (which cell to load?), and full cube results, see [`mjbommar/mimelens-001`](https://huggingface.co/mjbommar/mimelens-001).
+
+## How to use
+
+```python
+import torch
+from transformers import AutoModel
+from tokenizers import Tokenizer
+
+repo = "mjbommar/mimelens-001-tiny-bpe-64k-s2"
+model = AutoModel.from_pretrained(repo, trust_remote_code=True).eval()
+tok = Tokenizer.from_pretrained("mjbommar/binary-tokenizer-001-64k")
+cfg = model.config
+
+# BPE cell: encode 4 KB of raw bytes via the published binary-BPE tokenizer.
+window = open("path/to/file", "rb").read(4096)
+body = tok.encode(window.decode("latin-1")).ids[:1022]
+ids = [cfg.cls_token_id] + body + [cfg.sep_token_id]
+input_ids = torch.tensor([ids])
+
+with torch.no_grad():
+    out = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids))
+
+embedding = out.pooler_output     # (1, 256)  mean-pooled body-token embedding
+# Downstream: a frozen LR probe, a kNN over a labeled gallery, or fine-tune a classification head.
+# See the paper for the standard evaluation protocol.
+```
+
+## What this cell is
+
+- **Family**: [MimeLens-001](https://huggingface.co/mjbommar/mimelens-001) — 28 pretrained checkpoints across 3 sizes × 4 vocabularies × 2 seeds, plus one matched-tokens-seen ablation.
+- **Size**: `tiny` — 3.15 M backbone params, 4 layers, hidden 256, 4 attention heads, head dim 64.
+- **Input pipeline**: `bpe-64k` (65{,}536-entry binary BPE tokenizer (from binary-tokenizer-001-64k), ~2.09 bytes per token on the corpus.).
+- **Seed**: `2` (1 of 2 for this (size, vocab) combination).
+- **Pretraining**: 22,888 gradient updates, MLM-only, 30% mask ratio, 1024-token windows sampled uniformly at random across files and 64 KB fragments. AdamW + cosine LR (peak 5e-4, 2,000-step warmup, 10% floor), bf16 mixed precision, single RTX 4060 Ti.
+- **License**: MIT.
+
+## Evaluation
+
+Numbers below are for **this specific cell** on the `magic-frags` held-out test set (4 KB head of 64 KB random chunks, n=4,096). The within-cube comparison (3 sizes × 4 vocabs × 2-3 seeds, bootstrap CIs, adversarial sweep, calibration, real-network and disk-block validations) is in the [paper](https://github.com/mjbommar/binary-embedding-paper).
+
+| Benchmark | This cell |
+|---|---|
+| MIME-125 top-1 (magic-frags 4 KB head, n=4,096)        | **0.732** |
+| MIME-125 macro-F1 (magic-frags 4 KB head)              | 0.609 |
+| kNN R@1 (magic-frags, 3,147-file gallery / 949 queries) | 0.675 |
+
+## Recommended deployment regimes
+
+See the family hub README ([`mjbommar/mimelens-001`](https://huggingface.co/mjbommar/mimelens-001)) for the regime decision tree.
+
+## Training
+
+This cell is one point of the pre-registered 3 × 4 × 2 factorial cube described in the [MimeLens paper](https://github.com/mjbommar/binary-embedding-paper). Salient details:
+
+- **33 GB stratified multi-source binary corpus** (binary-30k + magic-frags + glaurung + Windows drivers).
+- **Position-arbitrary windowing**: 1024-token windows sampled uniformly at random across files and 64 KB fragments — no privileged "head of file" position. This is what makes MimeLens work on streaming / partial / random-offset inputs that whole-file detectors were not designed for.
+- **MLM-only** objective, 30% mask ratio (BERT replacement schedule: 80% `[MASK]`, 10% random, 10% original); tied input/output embeddings.
+- **Mean-pool over body tokens** for downstream tasks; the BERT-style `cls_pool` linear projection is *not* used because under MLM-only training it receives no gradient and remains at random init across all 28 cube cells (paper §3.4 verifies this).
+- **Wall-clock**: ~2.7 h on a single RTX 4060 Ti.
+
+## Honest caveats
+
+- This is one cell of a 28-cell cube. Within-cube comparisons in the paper come with bootstrap CIs at n=2 seeds; some marginal orderings (byte vs bpe-16k at the top of medium) are within seed noise and should be read as ties.
+- The training corpus is one 33 GB stratified multi-source binary sample. Results may not transfer to substantially different corpora.
+- All numbers are computed on data derived from a single labelling pipeline (libmagic-pinned via the [magic-bpe](https://github.com/mjbommar/magic-bpe) project). Cross-validation against PRONOM, Siegfried, DROID, or IANA reference files is a documented limitation.
+- CPU latency at the `medium` size is ~348× slower than Magika; for sub-millisecond whole-file triage on broad categories, Magika is purpose-built and is the right default. MimeLens occupies a different point on the deployment surface (position-arbitrary inputs + libmagic's 125-class taxonomy) rather than a drop-in replacement.
+- End-to-end fine-tuning on the production label distribution may shift these numbers and should be evaluated before deployment. The frozen-probe numbers reported above are not claimed as a lower bound on fine-tuned performance.
+
+## Citation
+
+```bibtex
+@misc{bommarito2026mimelens,
+  title  = {MimeLens: Pretrained Encoders for Fine-Grained Content-Type Detection},
+  author = {Bommarito II, Michael J.},
+  year   = {2026},
+  note   = {https://github.com/mjbommar/binary-embedding-paper},
+}
+```
+
+## Acknowledgments
+
+Thanks to the [magic-bpe](https://github.com/mjbommar/magic-bpe) project and the [binary-tokenizer-001](https://huggingface.co/mjbommar/binary-tokenizer-001-64k) family for the labelled corpus and BPE tokenizers this work builds on, and to the [Magika](https://github.com/google/magika) team for releasing a public package that made the §3 calibration possible.

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "MimeLensModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mimelens.MimeLensConfig",
+    "AutoModel": "modeling_mimelens.MimeLensModel"
+  },
+  "model_type": "mimelens",
+  "torch_dtype": "float32",
+  "vocab_size": 65543,
+  "hidden_size": 256,
+  "num_hidden_layers": 4,
+  "num_attention_heads": 4,
+  "head_dim": 64,
+  "ffn_multiplier_num": 8,
+  "ffn_multiplier_den": 3,
+  "max_position_embeddings": 1024,
+  "rope_theta": 10000.0,
+  "rms_norm_eps": 1e-06,
+  "pad_token_id": 2,
+  "cls_token_id": 4,
+  "sep_token_id": 5,
+  "mask_token_id": 6,
+  "byte_offset": 5,
+  "cls_pool_dim": 256,
+  "mimelens_cell_id": "tiny/bpe-64k/s2",
+  "mimelens_vocab_pipeline": "bpe-64k",
+  "mimelens_tokenizer_hub_id": "mjbommar/binary-tokenizer-001-64k",
+  "mimelens_pretraining_steps": 22888,
+  "mimelens_seed": 2
+}

configuration_mimelens.py

@@ -0,0 +1,126 @@
+"""HuggingFace-compatible config class for MimeLens.
+
+Copied verbatim into each per-cell HF repo (`mjbommar/mimelens-001-*`). Lets users
+do `AutoConfig.from_pretrained("mjbommar/mimelens-001-medium-bpe-16k-s1",
+trust_remote_code=True)` after the auto_map in config.json is honored.
+
+This file has zero non-stdlib dependencies beyond `transformers`. It's
+intentionally short — all torch / nn imports live in modeling_mimelens.py.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+
+class MimeLensConfig(PretrainedConfig):
+    """Configuration for a MimeLens encoder cell.
+
+    A MimeLens cell is one (size × vocab × seed) point of the binary-embedding
+    paper's 3x4xN cube: a BERT-style transformer encoder pretrained MLM-only on
+    33 GB of position-arbitrary binary content, with one of four input
+    pipelines (raw bytes, or BPE at 4K/16K/64K vocabulary).
+
+    For the architectural rationale and pretraining details see
+    docs/02-model-architecture.md and docs/04-training-protocol.md in the
+    paper repository (https://github.com/mjbommar/binary-embedding-paper).
+
+    Args:
+        vocab_size: int — full vocabulary including 5 special tokens. byte
+            cells: 261 (256 bytes + 5 specials). BPE cells: 4101 / 16391 / 65543.
+        hidden_size: int — transformer model dimension (256 / 384 / 512 for
+            tiny / small / medium).
+        num_hidden_layers: int — layer count (4 / 8 / 12 for tiny / small /
+            medium).
+        num_attention_heads: int — head count (4 / 6 / 8). Head dim is always
+            64 by design.
+        head_dim: int — per-head attention dimension. Fixed at 64 in the paper.
+        ffn_multiplier_num / ffn_multiplier_den: int — GeGLU FFN expansion as
+            a rational (8/3 ≈ 2.67, the LLaMA convention).
+        max_position_embeddings: int — RoPE position table size. Fixed at 1024
+            in the paper.
+        rope_theta: float — RoPE base frequency. Fixed at 10,000.
+        rms_norm_eps: float — RMSNorm epsilon. Fixed at 1e-6.
+        pad_token_id / cls_token_id / sep_token_id / mask_token_id: int —
+            special-token indices, matching binary_embedding.constants.
+        byte_offset: int — for byte cells, ord(b)+byte_offset gives the token
+            id. Fixed at 5 (after the 5 special tokens). Unused for BPE cells.
+        cls_pool_dim: int — output dim of the cls_pool layer. Note: this layer
+            receives no gradient under MLM-only training (see paper §3.4); the
+            mean-pool over body tokens is the trained pooling, not cls_pool.
+        initializer_range: float — std of trunc-normal init.
+
+        # MimeLens-specific provenance / tokenizer metadata:
+        mimelens_cell_id: str — e.g. "medium/bpe-16k/s1".
+        mimelens_vocab_pipeline: str — one of "byte", "bpe-4k", "bpe-16k",
+            "bpe-64k". Drives the tokenization in modeling_mimelens.
+        mimelens_tokenizer_hub_id: Optional[str] — for BPE cells, the HF Hub
+            id of the canonical tokenizer (e.g.
+            "mjbommar/binary-tokenizer-001-16k"). None for byte cells.
+        mimelens_pretraining_steps: int — total gradient updates (22,888
+            standard; 47,808 for the matched-tokens-seen ablation cell).
+        mimelens_seed: int — pretraining RNG seed (1, 2, or 3).
+    """
+
+    model_type = "mimelens"
+
+    def __init__(
+        self,
+        vocab_size: int = 16391,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 12,
+        num_attention_heads: int = 8,
+        head_dim: int = 64,
+        ffn_multiplier_num: int = 8,
+        ffn_multiplier_den: int = 3,
+        max_position_embeddings: int = 1024,
+        rope_theta: float = 10_000.0,
+        rms_norm_eps: float = 1e-6,
+        pad_token_id: int = 2,
+        cls_token_id: int = 4,
+        sep_token_id: int = 5,
+        mask_token_id: int = 6,
+        byte_offset: int = 5,
+        cls_pool_dim: int = 256,
+        initializer_range: float = 0.02,
+        mimelens_cell_id: str = "medium/bpe-16k/s1",
+        mimelens_vocab_pipeline: str = "bpe-16k",
+        mimelens_tokenizer_hub_id: Optional[str] = "mjbommar/binary-tokenizer-001-16k",
+        mimelens_pretraining_steps: int = 22_888,
+        mimelens_seed: int = 1,
+        **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.head_dim = head_dim
+        self.ffn_multiplier_num = ffn_multiplier_num
+        self.ffn_multiplier_den = ffn_multiplier_den
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.cls_token_id = cls_token_id
+        self.sep_token_id = sep_token_id
+        self.mask_token_id = mask_token_id
+        self.byte_offset = byte_offset
+        self.cls_pool_dim = cls_pool_dim
+        self.initializer_range = initializer_range
+        self.mimelens_cell_id = mimelens_cell_id
+        self.mimelens_vocab_pipeline = mimelens_vocab_pipeline
+        self.mimelens_tokenizer_hub_id = mimelens_tokenizer_hub_id
+        self.mimelens_pretraining_steps = mimelens_pretraining_steps
+        self.mimelens_seed = mimelens_seed
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+    @property
+    def head_size(self) -> int:
+        """For HF compatibility — alias for head_dim."""
+        return self.head_dim
+
+    @property
+    def intermediate_size(self) -> int:
+        """GeGLU expansion: hidden * (ffn_multiplier_num / ffn_multiplier_den)."""
+        return self.hidden_size * self.ffn_multiplier_num // self.ffn_multiplier_den

manifest.json

@@ -0,0 +1,11 @@
+{
+  "mimelens_release": "001",
+  "cell_id": "tiny/bpe-64k/s2",
+  "ckpt_source": "/data0/binary-embedding/phase-b/runs/tiny/bpe-64k/s2/checkpoints/best.safetensors",
+  "ckpt_sha256": "b2c5a5c935b300ad14eb2fd5efcb39dc2af784aa8fc40ca5ac7f3d8ed5b5d8ae",
+  "magicfiles_top1": 0.732421875,
+  "magicfiles_f1": 0.6086128821558999,
+  "magicfrags_top1": 0.732421875,
+  "magicfrags_f1": 0.6086128821558999,
+  "params_m": 3.15
+}

model.safetensors

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

modeling_mimelens.py

@@ -0,0 +1,274 @@
+"""HuggingFace-compatible inference model for MimeLens.
+
+Copied verbatim into each per-cell HF repo (`mjbommar/mimelens-001-*`). Lets
+users do:
+
+    from transformers import AutoModel, AutoConfig
+    config = AutoConfig.from_pretrained("mjbommar/mimelens-001-medium-bpe-16k-s1",
+                                         trust_remote_code=True)
+    model = AutoModel.from_pretrained("mjbommar/mimelens-001-medium-bpe-16k-s1",
+                                       trust_remote_code=True)
+    # → forward(input_ids, attention_mask) returns the mean-pooled body-token
+    #   embedding, shape (batch, hidden_size).
+
+The architecture is the small ModernBERT-style encoder from
+binary_embedding.models.encoder, vendored here to make each HF repo
+self-contained (no pip install binary_embedding required at inference time).
+
+Parameter naming is byte-compatible with the saved best.safetensors files so
+that AutoModel.from_pretrained() loads weights without prefix surgery.
+
+Pure torch; no scapy / sklearn / external deps. The mean-pool returned here
+is the same projection used throughout the paper; the cls_pool layer is
+known to receive no gradient under MLM-only training (see paper §3.4) and is
+kept only for state-dict compatibility — do not use it for downstream tasks.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+from .configuration_mimelens import MimeLensConfig
+
+
+# ---------------------------------------------------------------------------
+# Building blocks
+# ---------------------------------------------------------------------------
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm without bias. bf16-safe (norm computed in fp32)."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        variance = x.float().pow(2).mean(-1, keepdim=True)
+        normed = x * torch.rsqrt(variance + self.eps).to(x.dtype)
+        return normed * self.weight
+
+
+def _build_rope_cache(seq_len: int, head_dim: int, base: float,
+                      device: torch.device, dtype: torch.dtype):
+    positions = torch.arange(seq_len, device=device, dtype=torch.float32)
+    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device,
+                                            dtype=torch.float32) / head_dim))
+    freqs = torch.einsum("p,d->pd", positions, inv_freq)
+    cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).to(dtype)
+    sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).to(dtype)
+    return cos, sin
+
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply rotary position embedding. x: (..., seq, head_dim)."""
+    d = x.shape[-1]
+    x1, x2 = x[..., : d // 2], x[..., d // 2 :]
+    rotated = torch.cat([-x2, x1], dim=-1)
+    return x * cos + rotated * sin
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, head_dim: int):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.qkv = nn.Linear(hidden_size, 3 * num_heads * head_dim, bias=False)
+        self.out = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+                attn_mask: torch.Tensor) -> torch.Tensor:
+        B, S, _ = x.shape
+        qkv = self.qkv(x).reshape(B, S, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)  # each (B, S, H, D)
+        q = _apply_rope(q.transpose(1, 2), cos, sin)  # (B, H, S, D)
+        k = _apply_rope(k.transpose(1, 2), cos, sin)
+        v = v.transpose(1, 2)
+        # attn_mask: (B, 1, 1, S) with -inf at pad positions, 0 at real positions
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        out = out.transpose(1, 2).contiguous().reshape(B, S, -1)
+        return self.out(out)
+
+
+class FFN(nn.Module):
+    """GeGLU FFN: gelu(w_gate(x)) * w_up(x) → w_down."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.w_gate = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.w_up = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.w_down = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w_down(F.gelu(self.w_gate(x)) * self.w_up(x))
+
+
+class Layer(nn.Module):
+    def __init__(self, config: MimeLensConfig):
+        super().__init__()
+        self.norm1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = Attention(config.hidden_size, config.num_attention_heads,
+                              config.head_dim)
+        self.norm2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.ffn = FFN(config.hidden_size, config.intermediate_size)
+
+    def forward(self, x, cos, sin, attn_mask):
+        x = x + self.attn(self.norm1(x), cos, sin, attn_mask)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Top-level model
+# ---------------------------------------------------------------------------
+
+
+class MimeLensModel(PreTrainedModel):
+    """MimeLens encoder: bytes → mean-pooled embedding.
+
+    Use `forward(input_ids, attention_mask)` and consume the `pooler_output`
+    field (mean over body tokens, skipping the CLS / SEP / PAD positions).
+    Last-hidden-state is also returned as `last_hidden_state` if you want to
+    do your own pooling.
+    """
+
+    config_class = MimeLensConfig
+    base_model_prefix = "mimelens"
+
+    def __init__(self, config: MimeLensConfig):
+        super().__init__(config)
+        self.config = config
+
+        # Parameter naming MUST match best.safetensors: flat (no `encoder.` prefix).
+        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([Layer(config) for _ in range(config.num_hidden_layers)])
+        self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # cls_pool is kept ONLY for safetensors compatibility — receives no
+        # gradient under MLM-only training; use mean-pool instead.
+        self.cls_pool = nn.Linear(config.hidden_size, config.cls_pool_dim, bias=False)
+
+        # Lazily-built RoPE cache (one per device/dtype combination).
+        self._rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None
+        self._rope_cache_meta: Optional[tuple[torch.device, torch.dtype, int]] = None
+
+        # No weight init here — we always load_state_dict from a pretrained
+        # checkpoint via from_pretrained(). HF complains if we don't provide
+        # an init_weights; provide the no-op version.
+        self.post_init()
+
+    def _init_weights(self, module):
+        """No-op: we always load from a pretrained checkpoint."""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _get_rope(self, seq_len: int, device: torch.device, dtype: torch.dtype):
+        meta = (device, dtype, seq_len)
+        if self._rope_cache_meta != meta:
+            self._rope_cache = _build_rope_cache(seq_len, self.config.head_dim,
+                                                 self.config.rope_theta,
+                                                 device=device, dtype=dtype)
+            self._rope_cache_meta = meta
+        return self._rope_cache
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ):
+        B, S = input_ids.shape
+        x = self.embed(input_ids)  # (B, S, H)
+
+        # Build SDPA attention mask: (B, 1, 1, S) additive, -inf at pad.
+        if attention_mask is None:
+            attention_mask = torch.ones(B, S, device=input_ids.device, dtype=torch.long)
+        # Convert to additive: real (=1) → 0, pad (=0) → -inf
+        # SDPA expects mask broadcastable to (B, H, S, S)
+        attn_mask = attention_mask.to(x.dtype)
+        attn_mask = (1.0 - attn_mask).masked_fill((1.0 - attn_mask).bool(),
+                                                    torch.finfo(x.dtype).min)
+        # shape: (B, 1, 1, S) — broadcasts over heads and queries
+        attn_mask = attn_mask.view(B, 1, 1, S)
+
+        cos, sin = self._get_rope(S, device=x.device, dtype=x.dtype)
+
+        for layer in self.layers:
+            x = layer(x, cos, sin, attn_mask)
+
+        x = self.final_norm(x)
+        last_hidden_state = x
+
+        # Mean-pool over BODY tokens (skip CLS @ pos 0, SEP @ pos lens-1, PAD).
+        # attention_mask is (B, S) of {0,1}.
+        lens = attention_mask.sum(dim=1, keepdim=True)  # (B, 1)
+        positions = torch.arange(S, device=x.device).unsqueeze(0)  # (1, S)
+        body_mask = (positions >= 1) & (positions < (lens - 1))  # (B, S) bool
+        body_mask_f = body_mask.to(x.dtype).unsqueeze(-1)  # (B, S, 1)
+        pooled = (x * body_mask_f).sum(dim=1) / body_mask_f.sum(dim=1).clamp(min=1)
+        # shape: (B, H)
+
+        if not return_dict:
+            return (last_hidden_state, pooled)
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled,
+        )
+
+    # ---------------------------------------------------------------------
+    # Helper utilities for users (not part of the standard HF surface)
+    # ---------------------------------------------------------------------
+
+    def encode_bytes(
+        self,
+        byte_window: bytes,
+        tokenizer=None,
+        seq_len: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Convenience: encode one raw byte window into a mean-pooled embedding.
+
+        For byte cells (`config.mimelens_vocab_pipeline == 'byte'`), tokenizer
+        is ignored. For BPE cells, pass a BinaryTokenizer (from
+        `mjbommar/binary-tokenizer-001-*`) or any tokenizer with `.encode(bytes)
+        -> list[int]`.
+
+        Returns a (1, hidden_size) tensor on the same device as the model.
+        """
+        seq_len = seq_len or self.config.max_position_embeddings
+        body = seq_len - 2
+        cls_id = self.config.cls_token_id
+        sep_id = self.config.sep_token_id
+        pad_id = self.config.pad_token_id
+
+        if self.config.mimelens_vocab_pipeline == "byte":
+            ids = [b + self.config.byte_offset for b in byte_window[:body]]
+        else:
+            if tokenizer is None:
+                raise ValueError(
+                    f"BPE cell {self.config.mimelens_cell_id} requires a tokenizer; "
+                    f"load with e.g. `_native.BinaryTokenizer.from_file(...)` from "
+                    f"{self.config.mimelens_tokenizer_hub_id}"
+                )
+            ids = list(tokenizer.encode(byte_window))[:body]
+        out_ids = [cls_id, *ids, sep_id]
+        attn = [1] * len(out_ids) + [0] * (seq_len - len(out_ids))
+        out_ids = out_ids + [pad_id] * (seq_len - len(out_ids))
+
+        device = next(self.parameters()).device
+        input_ids = torch.tensor([out_ids], dtype=torch.long, device=device)
+        attention_mask = torch.tensor([attn], dtype=torch.long, device=device)
+        with torch.inference_mode():
+            return self(input_ids, attention_mask=attention_mask).pooler_output