File size: 5,954 Bytes
687ec72 | 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 | """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)
# mask=None (MLM, no padding) -> flash/mem-efficient SDPA, never a materialised LxL math path.
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(): # GPT-2 style scaled residual init
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) # no mask: windows have no padding
sel = labels != -100
hs = self.head_ln(F.gelu(self.head_dense(h[sel]))) # logits only for masked positions
logits = F.linear(hs, self.encoder.tok_emb.weight, self.head_bias) # tied
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] # [B,L]
h = self.encoder(ids, keep[:, None, None, :]) # bool mask: True = attend
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)
|