BERT-base sentence embeddings trained with in-batch contrastive learning (Multiple Negatives Ranking Loss) on MultiNLI entailment pairs.
1import torch
2import torch.nn.functional as F
3from transformers import AutoTokenizer, AutoModel
4
5def mean_pooling(last_hidden_state, attention_mask):
6 mask = attention_mask.unsqueeze(-1).to(dtype=last_hidden_state.dtype)
7 summed = (last_hidden_state * mask).sum(dim=1)
8 counts = mask.sum(dim=1).clamp(min=1e-6)
9 return summed / counts
10
11@torch.no_grad()
12def embed_texts(texts, model_id="rafidka/vectra", max_length=128, device="cuda"):
13 tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
14 model = AutoModel.from_pretrained(model_id, add_pooling_layer=False).to(device).eval()
15 batch = tok(texts, padding="max_length", truncation=True, max_length=max_length, return_tensors="pt").to(device)
16 out = model(**batch)
17 emb = mean_pooling(out.last_hidden_state, batch["attention_mask"])
18 emb = F.normalize(emb, p=2, dim=-1)
19 return emb
