Spaces:
Sleeping
Sleeping
File size: 4,451 Bytes
4258647 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | # 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
# -----------------------
@property
def model_id(self) -> str:
return self._local_model_path or self._repo_id
@property
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 {}
|