File size: 1,922 Bytes
701cf7d | 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 | """
Vector embedding for RAG - uses model's own hidden states as embeddings (no external API)
Flaw fix: embedding drift -> L2 normalize, mean pool last hidden.
"""
import torch
import numpy as np
from typing import List
class AresEmbedder:
def __init__(self, model, tokenizer, device="cpu"):
self.model = model
self.tokenizer = tokenizer
self.device = device
self.model.eval()
@torch.no_grad()
def embed(self, texts: List[str], batch_size=8, max_len=512) -> np.ndarray:
embeddings = []
for i in range(0, len(texts), batch_size):
batch = texts[i:i+batch_size]
ids_batch = [self.tokenizer.encode(t)[:max_len] for t in batch]
max_l = max(len(x) for x in ids_batch)
padded = [x + [self.tokenizer.pad_token_id]*(max_l-len(x)) for x in ids_batch]
input_ids = torch.tensor(padded, dtype=torch.long, device=self.device)
out = self.model(input_ids=input_ids)
hidden = out["hidden_states"] # [b,s,hidden]
# Mean pooling ignoring pad
mask = (input_ids != self.tokenizer.pad_token_id).float().unsqueeze(-1) # [b,s,1]
summed = (hidden * mask).sum(dim=1)
counts = mask.sum(dim=1).clamp(min=1)
mean = summed / counts
# L2 normalize
norm = torch.nn.functional.normalize(mean, p=2, dim=1)
embeddings.append(norm.cpu().numpy())
if embeddings:
return np.concatenate(embeddings, axis=0)
return np.zeros((0, self.model.config.hidden_size))
@staticmethod
def cosine_similarity(query_emb: np.ndarray, doc_embs: np.ndarray) -> np.ndarray:
# query_emb: [hidden] or [1, hidden], doc_embs: [n, hidden]
if query_emb.ndim == 1:
query_emb = query_emb[None, :]
# Assume normalized
return np.dot(doc_embs, query_emb.T).squeeze()
|