File size: 5,819 Bytes
00a471f 3f8b024 00a471f 3f8b024 00a471f 3f8b024 00a471f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """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 7 special tokens (start, end,
pad, unk, cls, sep, mask). byte cells: 263 (256 bytes + 7 specials).
BPE cells: 4103 / 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 7 (after the 7 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 = 7,
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
|