File size: 1,111 Bytes
138e5ba | 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 | # e5_encoder.py
from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np
class E5Encoder:
def __init__(self, model_name: str = "intfloat/e5-base", device: str = None):
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name).to(self.device)
def encode(self, text: str, normalize: bool = True) -> np.ndarray:
if not text:
return np.zeros((self.model.config.hidden_size,), dtype=np.float32)
# Prefix for E5 models to work properly
prefixed_text = f"passage: {text}" if not text.startswith("query:") else text
inputs = self.tokenizer(prefixed_text, return_tensors="pt", truncation=True, padding=True).to(self.device)
with torch.no_grad():
output = self.model(**inputs)
emb = output.last_hidden_state[:, 0] # [CLS] token
emb = emb.squeeze(0)
if normalize:
emb = emb / torch.norm(emb, p=2)
return emb.cpu().numpy()
|