GigaProjects

← Back to rag-assistant

embeddings.py

"""Embedding helpers using sentence-transformers."""

from __future__ import annotations

import os
from functools import lru_cache

from src.config import get_embedding_model_name


os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


@lru_cache(maxsize=1)
def load_embedding_model():
    from huggingface_hub.utils import disable_progress_bars
    from sentence_transformers import SentenceTransformer

    disable_progress_bars()
    model_name = get_embedding_model_name()

    try:
        return SentenceTransformer(model_name, local_files_only=True)
    except Exception:
        return SentenceTransformer(model_name)


def embed_texts(texts: list[str]) -> list[list[float]]:
    if not texts:
        return []

    model = load_embedding_model()
    embeddings = model.encode(texts, normalize_embeddings=True)
    return embeddings.tolist()


def embed_query(question: str) -> list[float]:
    embeddings = embed_texts([question])
    return embeddings[0]

Run this code