| """Encoder-only embedding model. Pre-norm, GELU, learned positions, SDPA attention, tied MLM head.""" |
| import json, math |
| import torch, torch.nn as nn, torch.nn.functional as F |
| from torch.utils.checkpoint import checkpoint |
|
|
| PAD_ID, MASK_ID, VOCAB = 50256, 50257, 50258 |
|
|
| CFG_57M = dict(vocab=VOCAB, hidden=512, heads=8, ffn=2048, layers=10, max_pos=1024, dropout=0.0) |
|
|
| class Block(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| h = c["hidden"]; self.nh = c["heads"]; self.p = c["dropout"] |
| self.ln1 = nn.LayerNorm(h); self.qkv = nn.Linear(h, 3 * h); self.proj = nn.Linear(h, h) |
| self.ln2 = nn.LayerNorm(h); self.fc1 = nn.Linear(h, c["ffn"]); self.fc2 = nn.Linear(c["ffn"], h) |
| def forward(self, x, mask=None): |
| B, L, H = x.shape |
| q, k, v = self.qkv(self.ln1(x)).view(B, L, 3, self.nh, H // self.nh).permute(2, 0, 3, 1, 4) |
| |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.p if self.training else 0.0) |
| x = x + F.dropout(self.proj(a.transpose(1, 2).reshape(B, L, H)), self.p, self.training) |
| return x + F.dropout(self.fc2(F.gelu(self.fc1(self.ln2(x)))), self.p, self.training) |
|
|
| class Encoder(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.c = c |
| self.tok_emb = nn.Embedding(c["vocab"], c["hidden"]) |
| self.pos_emb = nn.Embedding(c["max_pos"], c["hidden"]) |
| self.blocks = nn.ModuleList([Block(c) for _ in range(c["layers"])]) |
| self.ln_f = nn.LayerNorm(c["hidden"]) |
| self.grad_ckpt = False |
| self.apply(self._init) |
| for n, p in self.named_parameters(): |
| if n.endswith("proj.weight") or n.endswith("fc2.weight"): |
| nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * c["layers"])) |
| @staticmethod |
| def _init(m): |
| if isinstance(m, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(m.weight, 0.0, 0.02) |
| if getattr(m, "bias", None) is not None: nn.init.zeros_(m.bias) |
| def forward(self, ids, attn_mask=None): |
| B, L = ids.shape |
| x = self.tok_emb(ids) + self.pos_emb(torch.arange(L, device=ids.device))[None] |
| for blk in self.blocks: |
| if self.grad_ckpt and self.training and torch.is_grad_enabled(): |
| x = checkpoint(blk, x, attn_mask, use_reentrant=False) |
| else: |
| x = blk(x, attn_mask) |
| return self.ln_f(x) |
|
|
| class MLMModel(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.encoder = Encoder(c) |
| h = c["hidden"] |
| self.head_dense = nn.Linear(h, h); self.head_ln = nn.LayerNorm(h) |
| self.head_bias = nn.Parameter(torch.zeros(c["vocab"])) |
| def forward(self, ids, labels): |
| h = self.encoder(ids) |
| sel = labels != -100 |
| hs = self.head_ln(F.gelu(self.head_dense(h[sel]))) |
| logits = F.linear(hs, self.encoder.tok_emb.weight, self.head_bias) |
| return F.cross_entropy(logits.float(), labels[sel]) |
|
|
| class EmbedModel(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.encoder = Encoder(c) |
| def forward(self, ids, lens): |
| L = ids.shape[1] |
| keep = torch.arange(L, device=ids.device)[None] < lens[:, None] |
| h = self.encoder(ids, keep[:, None, None, :]) |
| m = keep.unsqueeze(-1).to(h.dtype) |
| e = (h * m).sum(1) / m.sum(1).clamp_min(1) |
| return F.normalize(e.float(), dim=-1) |
|
|
| def count_params(model): |
| return sum(p.numel() for p in model.parameters()) |
|
|
| def mask_tokens(ids, p=0.15, gen=None): |
| """Dynamic BERT masking 80/10/10 on GPU. EOT (pad/doc separator) is never selected.""" |
| sel = (torch.rand(ids.shape, device=ids.device, generator=gen) < p) & (ids != PAD_ID) |
| labels = torch.where(sel, ids, torch.full_like(ids, -100)) |
| r = torch.rand(ids.shape, device=ids.device, generator=gen) |
| inp = ids.clone() |
| inp[sel & (r < 0.8)] = MASK_ID |
| rnd = sel & (r >= 0.8) & (r < 0.9) |
| inp[rnd] = torch.randint(0, PAD_ID, ids.shape, device=ids.device, generator=gen)[rnd] |
| return inp, labels |
|
|
| def load_encoder_weights(model, path, map_location="cpu"): |
| sd = torch.load(path, map_location=map_location) |
| sd = sd.get("model", sd) |
| enc = {k[len("encoder."):]: v for k, v in sd.items() if k.startswith("encoder.")} or sd |
| missing, unexpected = model.encoder.load_state_dict(enc, strict=True), None |
| return missing |
|
|
| class Embedder: |
| """Usage: e = Embedder('path/to/57M_embed_Full'); e.encode(['hello world'])""" |
| def __init__(self, folder, device=None): |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
| self.c = json.load(open(f"{folder}/config.json"))["encoder_config"] |
| self.m = EmbedModel(self.c); self.m.encoder.load_state_dict(load_file(f"{folder}/model.safetensors")) |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.m.to(self.device).eval() |
| self.tok = AutoTokenizer.from_pretrained("openai-community/gpt2") |
| @torch.no_grad() |
| def encode(self, texts, max_len=128, bs=64): |
| out = [] |
| for i in range(0, len(texts), bs): |
| ids = self.tok(texts[i:i+bs], truncation=True, max_length=max_len, add_special_tokens=False)["input_ids"] |
| L = max(1, max(len(x) for x in ids)) |
| t = torch.full((len(ids), L), PAD_ID, dtype=torch.long) |
| lens = torch.tensor([max(1, len(x)) for x in ids]) |
| for j, x in enumerate(ids): t[j, :len(x)] = torch.tensor(x, dtype=torch.long) |
| out.append(self.m(t.to(self.device), lens.to(self.device)).cpu()) |
| return torch.cat(out) |
|
|
