Spaces:
Running
Running
| from __future__ import annotations | |
| import os | |
| import time | |
| from typing import Iterable | |
| from dotenv import load_dotenv | |
| from huggingface_hub import InferenceClient | |
| from langchain_core.embeddings import Embeddings | |
| import numpy as np | |
| from sklearn.feature_extraction.text import HashingVectorizer | |
| DEFAULT_EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2" | |
| DEFAULT_EMBED_PROVIDER = "local" | |
| DEFAULT_EMBED_BATCH_SIZE = 16 | |
| DEFAULT_EMBED_RETRIES = 3 | |
| load_dotenv() | |
| def get_hf_token() -> str | None: | |
| return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN") | |
| def to_single_vector(result) -> list[float]: | |
| array = np.asarray(result, dtype=np.float32) | |
| if array.ndim == 1: | |
| vector = array | |
| elif array.ndim == 2: | |
| vector = array.mean(axis=0) | |
| elif array.ndim == 3: | |
| vector = array[0].mean(axis=0) | |
| else: | |
| raise RuntimeError(f"Unexpected embedding shape: {array.shape}") | |
| norm = np.linalg.norm(vector) | |
| if norm: | |
| vector = vector / norm | |
| return vector.astype(float).tolist() | |
| class HFTextEmbeddings(Embeddings): | |
| def __init__(self, model: str | None = None, provider: str | None = None) -> None: | |
| self.model = model or os.getenv("HF_EMBED_MODEL", DEFAULT_EMBED_MODEL) | |
| self.provider = provider or os.getenv("HF_EMBED_PROVIDER", DEFAULT_EMBED_PROVIDER) | |
| self.batch_size = int(os.getenv("HF_EMBED_BATCH_SIZE", str(DEFAULT_EMBED_BATCH_SIZE))) | |
| self.retries = int(os.getenv("HF_EMBED_RETRIES", str(DEFAULT_EMBED_RETRIES))) | |
| self.client = None | |
| self.local_vectorizer = None | |
| if self.provider == "local": | |
| # A stateless vectorizer keeps indexing and querying compatible without | |
| # downloading a model or consuming Hugging Face inference credits. | |
| self.local_vectorizer = HashingVectorizer( | |
| n_features=384, | |
| analyzer="char_wb", | |
| ngram_range=(3, 5), | |
| lowercase=True, | |
| alternate_sign=False, | |
| norm="l2", | |
| ) | |
| else: | |
| self.client = InferenceClient(provider=self.provider, api_key=get_hf_token()) | |
| def embed_documents(self, texts: list[str]) -> list[list[float]]: | |
| return self._embed(texts) | |
| def embed_query(self, text: str) -> list[float]: | |
| return self._embed([text])[0] | |
| def _embed(self, texts: Iterable[str]) -> list[list[float]]: | |
| inputs = list(texts) | |
| if not inputs: | |
| return [] | |
| if self.local_vectorizer is not None: | |
| print(f"Local embeddings payload: inputs={len(inputs)} dimensions=384", flush=True) | |
| return self.local_vectorizer.transform(inputs).toarray().astype(float).tolist() | |
| vectors: list[list[float]] = [] | |
| for start in range(0, len(inputs), self.batch_size): | |
| batch = inputs[start : start + self.batch_size] | |
| print( | |
| f"Outgoing HF embeddings payload: model={self.model} inputs={len(batch)} " | |
| f"offset={start}", | |
| flush=True, | |
| ) | |
| for text in batch: | |
| vectors.append(self._embed_one(text)) | |
| return vectors | |
| def _embed_one(self, text: str) -> list[float]: | |
| last_error: Exception | None = None | |
| for attempt in range(1, self.retries + 1): | |
| try: | |
| return to_single_vector( | |
| self.client.feature_extraction( | |
| text, | |
| model=self.model, | |
| normalize=True, | |
| truncate=True, | |
| ) | |
| ) | |
| except Exception as exc: | |
| last_error = exc | |
| if attempt >= self.retries: | |
| break | |
| time.sleep(min(2 * attempt, 5)) | |
| raise RuntimeError(f"HF embeddings failed after {self.retries} attempts: {last_error}") from last_error | |