spb/uqo-research-db

1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819import numpy as np2021from src.utils.config import get_embeddings_config22from src.utils.logging import get_logger2324logger = get_logger("embeddings.models")252627class EmbeddingModelRegistry:2829    def __init__(self) -> None:30        self._config = get_embeddings_config()31        self._models: dict[str, dict] = self._config.get("models", {})32        self._default_model: str = self._config.get("default_model", "all-MiniLM-L6-v2")33        self._loaded: dict = {}3435    @property36    def default_model(self) -> str:37        return self._default_model3839    def list_models(self) -> list[dict]:40        result: list[dict] = []41        for name, info in self._models.items():42            entry = {"name": name, **info}43            entry["is_default"] = name == self._default_model44            result.append(entry)45        return result4647    def get_model_info(self, model_name: str) -> dict:48        if model_name not in self._models:49            available = ", ".join(self._models.keys())50            raise ValueError(51                f"Unknown model '{model_name}'. Available: {available}"52            )53        info = dict(self._models[model_name])54        info["name"] = model_name55        return info5657    def load_model(self, model_name: str):58        if model_name in self._loaded:59            return self._loaded[model_name]6061        info = self.get_model_info(model_name)62        provider = info.get("provider", "")6364        if provider == "openai":65            raise NotImplementedError(66                f"OpenAI model '{model_name}' requires API key configuration. "67                "Set OPENAI_API_KEY environment variable and implement the OpenAI "68                "embedding client to use this model."69            )7071        if provider != "sentence-transformers":72            raise ValueError(f"Unsupported provider '{provider}' for model '{model_name}'")7374        from sentence_transformers import SentenceTransformer7576        device = self._config.get("processing", {}).get("device", "cpu")77        logger.info("Loading sentence-transformers model '%s' on device '%s'", model_name, device)78        model = SentenceTransformer(model_name, device=device)79        self._loaded[model_name] = model80        logger.info("Model '%s' loaded successfully", model_name)81        return model8283    def encode(84        self,85        texts: list[str],86        model_name: str | None = None,87        batch_size: int = 256,88        show_progress: bool = True,89    ) -> np.ndarray:90        model_name = model_name or self._default_model91        info = self.get_model_info(model_name)92        provider = info.get("provider", "")9394        if provider == "openai":95            raise NotImplementedError(96                f"OpenAI model '{model_name}' requires API key configuration. "97                "Set OPENAI_API_KEY environment variable and implement the OpenAI "98                "embedding client to use this model."99            )100101        model = self.load_model(model_name)102        normalize = self._config.get("processing", {}).get("normalize", True)103104        logger.info(105            "Encoding %d texts with '%s' (batch_size=%d)",106            len(texts), model_name, batch_size,107        )108        vectors = model.encode(109            texts,110            batch_size=batch_size,111            show_progress_bar=show_progress,112            normalize_embeddings=normalize,113        )114        return np.asarray(vectors)115