File size: 1,922 Bytes
78b2584 | 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 | #!/usr/bin/env python3
"""
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) # attention pooling
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() |