Views
No views yet
1import torch
2import numpy as np
3from transformers import AutoTokenizer, AutoModel
4
5class Encoder:
6
7 def __init__(self, model_name, device="cuda", max_length=32):
8
9 self.device = device
10 self.max_length = max_length
11
12 self.tokenizer = AutoTokenizer.from_pretrained(model_name)
13 self.model = AutoModel.from_pretrained(model_name)
14
15 if device == "cuda":
16 self.model = self.model.cuda()
17
18 self.model.eval()
19
20 def encode(self, texts, batch_size=256):
21
22 all_vecs = []
23
24 with torch.no_grad():
25 for i in range(0, len(texts), batch_size):
26
27 batch = texts[i:i+batch_size]
28
29 tokens = self.tokenizer(
30 batch,
31 padding=True,
32 truncation=True,
33 max_length=self.max_length,
34 return_tensors="pt"
35 )
36
37 if self.device == "cuda":
38 tokens = {k: v.cuda() for k, v in tokens.items()}
39
40 out = self.model(**tokens)
41
42 hidden = out.last_hidden_state
43 mask = tokens["attention_mask"].unsqueeze(-1)
44
45 pooled = (hidden * mask).sum(1) / mask.sum(1)
46
47 # IMPORTANT: normalize embeddings
48 pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
49
50 all_vecs.append(pooled.cpu().numpy())
51
52 return np.vstack(all_vecs).astype("float32")