File size: 1,827 Bytes
feaa032 | 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 | import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
from config import EMBEDDING_MODEL, EMBEDDING_DIM, resolve_device
class HypothesisEncoder:
def __init__(
self,
model_name=EMBEDDING_MODEL,
device=None
):
self.device = device or resolve_device()
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name).to(self.device)
self.model.eval()
for p in self.model.parameters():
p.requires_grad = False
self.dim = self.model.config.hidden_size
if model_name == EMBEDDING_MODEL:
# Everything downstream (dataset.py, sample_inference.py, ...)
# reads dim from config.EMBEDDING_DIM rather than this instance,
# so a silent mismatch here would surface as a shape error far
# away from its cause. Catch it at the source instead.
assert self.dim == EMBEDDING_DIM, (
f"{model_name} produces {self.dim}-dim embeddings but "
f"config.EMBEDDING_DIM is {EMBEDDING_DIM}. Update config.py."
)
print(f"[ENCODER] Using {model_name}")
print(f"[ENCODER] Embedding dimension: {self.dim}")
@torch.no_grad()
def encode(self, texts):
inputs = self.tokenizer(
texts, padding=True, truncation=True, max_length=128, return_tensors="pt"
).to(self.device)
outputs = self.model(**inputs)
hidden = outputs.last_hidden_state
mask = inputs["attention_mask"].unsqueeze(-1)
pooled = (hidden * mask).sum(dim=1)
counts = mask.sum(dim=1)
pooled = pooled / counts.clamp(min=1e-9)
pooled = F.normalize(pooled, dim=-1)
return pooled
|