Spaces:
Sleeping
Sleeping
| # get_embedding.py | |
| from __future__ import annotations | |
| import asyncio | |
| import os | |
| import yaml | |
| import torch | |
| from typing import List, Optional | |
| from huggingface_hub import snapshot_download | |
| from langchain_huggingface import HuggingFaceEmbeddings | |
| # If you have an auth helper, we keep it | |
| from login import HuggingFaceLogin | |
| class EmbeddingFetcher: | |
| """ | |
| Async-friendly wrapper around a HuggingFace embedding model. | |
| - Lazily initializes the model and downloads repo snapshot at app startup. | |
| - Runs blocking HF / Torch operations in a worker thread via asyncio.to_thread. | |
| - Thread-safe through an asyncio.Lock to prevent duplicate initialization. | |
| """ | |
| def __init__(self, config_path: str = "config.yml") -> None: | |
| self._config_path = config_path | |
| self._ready = False | |
| self._init_lock = asyncio.Lock() | |
| self._model: Optional[HuggingFaceEmbeddings] = None | |
| self._local_model_path: Optional[str] = None | |
| # Authenticate early (likely blocking), but off main thread in ensure_ready() | |
| self._login = HuggingFaceLogin() | |
| # Config defaults (overridden by config.yml) | |
| self._repo_id: str = "SPAL0028/default-model" | |
| self._normalize_embeddings: bool = False | |
| # Device | |
| self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # ----------------------- | |
| # Public properties | |
| # ----------------------- | |
| def model_id(self) -> str: | |
| return self._local_model_path or self._repo_id | |
| def device_str(self) -> str: | |
| return str(self._device) | |
| # ----------------------- | |
| # Initialization | |
| # ----------------------- | |
| async def ensure_ready(self) -> None: | |
| """ | |
| Idempotent async initializer. Safe to call multiple times. | |
| """ | |
| if self._ready: | |
| return | |
| async with self._init_lock: | |
| if self._ready: | |
| return | |
| # 1) Login (may prompt environment-based auth); keep this off the main loop | |
| await asyncio.to_thread(self._login.authenticate) | |
| # 2) Load config | |
| cfg = await asyncio.to_thread(self._load_config) | |
| self._repo_id = cfg.get("hugging_face_url", self._repo_id) | |
| self._normalize_embeddings = bool(cfg.get("normalize_embeddings", self._normalize_embeddings)) | |
| # 3) Download snapshot if needed (blocking IO) off main thread | |
| self._local_model_path = await asyncio.to_thread( | |
| snapshot_download, | |
| self._repo_id | |
| ) | |
| # 4) Build embeddings (blocking CPU-bound creation) off main thread | |
| def _build_embeddings(): | |
| return HuggingFaceEmbeddings( | |
| model_name=self._local_model_path, | |
| model_kwargs={"device": self._device}, | |
| encode_kwargs={"normalize_embeddings": self._normalize_embeddings}, | |
| ) | |
| self._model = await asyncio.to_thread(_build_embeddings) | |
| self._ready = True | |
| # ----------------------- | |
| # Core API | |
| # ----------------------- | |
| async def embed(self, texts: List[str] | str) -> List[List[float]]: | |
| """ | |
| Generate embeddings for a list of texts. | |
| Ensures the model is initialized and runs blocking ops using a thread. | |
| """ | |
| await self.ensure_ready() | |
| if isinstance(texts, str): | |
| texts = [texts] | |
| if not texts: | |
| raise ValueError("No texts provided for embedding.") | |
| # Defensive strip to avoid empty items sneaking through | |
| sanitized = [t if isinstance(t, str) else str(t) for t in texts] | |
| sanitized = [t.strip() for t in sanitized if t and t.strip()] | |
| if not sanitized: | |
| raise ValueError("All provided texts are empty after sanitization.") | |
| # langchain_huggingface.HuggingFaceEmbeddings.embed_documents is blocking | |
| vectors = await asyncio.to_thread(self._model.embed_documents, sanitized) # type: ignore[union-attr] | |
| return vectors | |
| # ----------------------- | |
| # Helpers | |
| # ----------------------- | |
| def _load_config(self) -> dict: | |
| if not os.path.exists(self._config_path): | |
| # Keep behavior predictable if config is missing | |
| return {} | |
| with open(self._config_path, "r", encoding="utf-8") as f: | |
| return yaml.safe_load(f) or {} | |