| |
| """ |
| Shared encoder used by 13 (build index) and 14 (profile/eval). |
| |
| The layer names here match what 12_train_encoder.py saved, so a checkpoint trained by |
| 12 loads straight into this class. |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import AutoModel, AutoTokenizer |
|
|
|
|
| class Encoder(nn.Module): |
| def __init__(self, model_name, emb_dim): |
| super().__init__() |
| self.backbone = AutoModel.from_pretrained(model_name, trust_remote_code=True) |
| h = self.backbone.config.hidden_size |
| self.attn = nn.Linear(h, 1) |
| self.proj = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, emb_dim)) |
|
|
| def forward(self, input_ids, attention_mask): |
| hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask)[0] |
| scores = self.attn(hidden).squeeze(-1).masked_fill(attention_mask == 0, -1e9) |
| w = scores.softmax(-1).unsqueeze(-1) |
| pooled = (hidden * w).sum(1) |
| return F.normalize(self.proj(pooled), dim=-1) |
|
|
|
|
| def load_encoder(ckpt_path, device): |
| ck = torch.load(ckpt_path, map_location=device) |
| cfg = ck["config"] |
| model = Encoder(cfg["MODEL"], cfg["EMB_DIM"]).to(device) |
| model.load_state_dict(ck["model"]) |
| model.eval() |
| tok = AutoTokenizer.from_pretrained(cfg["MODEL"], trust_remote_code=True) |
| return model, tok, cfg |
|
|
|
|
| @torch.no_grad() |
| def embed(model, tok, seqs, device, max_length=128, bf16=True): |
| """Return an (N, emb_dim) float32 numpy array of L2-normalized embeddings.""" |
| t = tok(seqs, padding=True, truncation=True, max_length=max_length, return_tensors="pt") |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16, |
| enabled=(bf16 and device == "cuda")): |
| z = model(t["input_ids"].to(device), t["attention_mask"].to(device)) |
| return z.float().cpu().numpy() |
