Spaces:
Running
Running
File size: 3,937 Bytes
9a1014e | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | 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
|