Spaces:
Running
Running
File size: 3,458 Bytes
ca58521 3adb6fd ca58521 3adb6fd ca58521 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """CPU-only, unnormalized embeddings from multilingual-e5-small."""
from __future__ import annotations
import math
import os
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Literal
MODEL_ID = "intfloat/multilingual-e5-small"
MODEL_DIR = Path(os.getenv("E5_MODEL_DIR", "/opt/models/multilingual-e5-small"))
EXPECTED_DIMENSIONS = 384
MAX_INPUT_TOKENS = 512
class TooManyTokens(ValueError):
def __init__(self, count: int, limit: int, input_index: int | None = None) -> None:
self.count = count
self.limit = limit
self.input_index = input_index
item = f"Input item {input_index}" if input_index is not None else "Input"
super().__init__(f"{item} contains {count} tokens; maximum is {limit} (including prefix and special tokens)")
@dataclass(frozen=True)
class EmbeddingResult:
tokens_count: int
embedding: list[float]
class EmbeddingEngine:
"""Run official E5 attention-mask-aware mean pooling without L2 normalization."""
def __init__(self, tokenizer: object, model: object, torch_module: object) -> None:
self.tokenizer = tokenizer
self.model = model
self.torch = torch_module
self._lock = threading.Lock() # Prevent multiple CPU-heavy inferences in parallel.
model_limit = int(getattr(model.config, "max_position_embeddings", MAX_INPUT_TOKENS))
self.max_tokens = min(MAX_INPUT_TOKENS, model_limit)
self.dimensions = int(model.config.hidden_size)
if self.dimensions != EXPECTED_DIMENSIONS:
raise RuntimeError(f"Unexpected embedding dimensions: {self.dimensions}")
@classmethod
def load(cls) -> EmbeddingEngine:
import torch
from transformers import AutoModel, AutoTokenizer
threads = max(1, min(2, int(os.getenv("CPU_THREADS", "2"))))
torch.set_num_threads(threads)
tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR), local_files_only=True)
model = AutoModel.from_pretrained(str(MODEL_DIR), local_files_only=True)
model.to("cpu")
model.eval() # Turn off dropout for repeatable output.
return cls(tokenizer, model, torch)
def encode(self, text: str, kind: Literal["query", "passage"]) -> EmbeddingResult:
# The prefix is required by E5. The caller's original text is never modified.
input_text = f"{kind}: {text}"
with self._lock:
encoded = self.tokenizer(
input_text,
return_tensors="pt",
add_special_tokens=True,
truncation=False,
)
tokens_count = int(encoded["attention_mask"].sum().item())
if tokens_count > self.max_tokens:
raise TooManyTokens(tokens_count, self.max_tokens)
with self.torch.inference_mode():
output = self.model(**encoded)
token_embeddings = output.last_hidden_state
attention_mask = encoded["attention_mask"].unsqueeze(-1)
pooled = (token_embeddings * attention_mask).sum(dim=1) / attention_mask.sum(dim=1).clamp(min=1)
vector = pooled.squeeze(0).cpu().tolist()
if len(vector) != self.dimensions or any(not math.isfinite(value) for value in vector):
raise RuntimeError("Model produced an invalid embedding")
return EmbeddingResult(tokens_count=tokens_count, embedding=vector)
|