dp-zp-agent/scripts/embedding_utils.py

from __future__ import annotations

import os
from functools import lru_cache
from typing import Any

import numpy as np


DEFAULT_EMBEDDING_MODEL = (
    "intfloat/multilingual-e5-small"
)

DEFAULT_BATCH_SIZE = 32


def embedding_model_name() -> str:
    return (
        os.getenv(
            "EMBEDDING_MODEL",
            DEFAULT_EMBEDDING_MODEL,
        ).strip()
        or DEFAULT_EMBEDDING_MODEL
    )


def embedding_batch_size() -> int:
    raw_value = os.getenv(
        "EMBEDDING_BATCH_SIZE",
        str(DEFAULT_BATCH_SIZE),
    ).strip()

    try:
        value = int(
            raw_value
        )

    except ValueError:
        return DEFAULT_BATCH_SIZE

    return max(
        1,
        value,
    )


@lru_cache(maxsize=2)
def get_embedding_model(
    model_name: str,
):
    from sentence_transformers import (
        SentenceTransformer,
    )

    return SentenceTransformer(
        model_name
    )


def build_embedding_text(
    chunk: dict[str, Any],
) -> str:
    parts: list[str] = []

    title = str(
        chunk.get("title")
        or ""
    ).strip()

    author = str(
        chunk.get("author")
        or ""
    ).strip()

    tags = [
        str(value).strip()
        for value in (
            chunk.get("tags")
            or []
        )
        if str(value).strip()
    ]

    categories = [
        str(value).strip()
        for value in (
            chunk.get("categories")
            or []
        )
        if str(value).strip()
    ]

    text = str(
        chunk.get("text")
        or ""
    ).strip()

    if title:
        parts.append(
            f"Názov: {title}"
        )

    if author:
        parts.append(
            f"Autor: {author}"
        )

    if tags:
        parts.append(
            "Tagy: "
            + ", ".join(
                tags
            )
        )

    if categories:
        parts.append(
            "Kategórie: "
            + ", ".join(
                categories
            )
        )

    if text:
        parts.append(
            text
        )

    return "\n".join(
        parts
    )


def embed_passages(
    texts: list[str],
    *,
    model_name: str | None = None,
) -> np.ndarray:
    if not texts:
        return np.empty(
            (
                0,
                0,
            ),
            dtype=np.float32,
        )

    selected_model = (
        model_name
        or embedding_model_name()
    )

    model = get_embedding_model(
        selected_model
    )

    prepared = [
        "passage: "
        + text.strip()
        for text in texts
    ]

    vectors = model.encode(
        prepared,
        batch_size=(
            embedding_batch_size()
        ),
        show_progress_bar=False,
        convert_to_numpy=True,
        normalize_embeddings=True,
    )

    return np.asarray(
        vectors,
        dtype=np.float32,
    )


def embed_query(
    query: str,
    *,
    model_name: str | None = None,
) -> np.ndarray:
    clean_query = query.strip()

    if not clean_query:
        raise ValueError(
            "Query nesmie byť prázdny"
        )

    selected_model = (
        model_name
        or embedding_model_name()
    )

    model = get_embedding_model(
        selected_model
    )

    vector = model.encode(
        [
            "query: "
            + clean_query
        ],
        batch_size=1,
        show_progress_bar=False,
        convert_to_numpy=True,
        normalize_embeddings=True,
    )[0]

    return np.asarray(
        vector,
        dtype=np.float32,
    )


def vector_to_blob(
    vector: np.ndarray,
) -> bytes:
    normalized = np.asarray(
        vector,
        dtype=np.float32,
    )

    return normalized.tobytes()


def blob_to_vector(
    blob: bytes,
    dimensions: int,
) -> np.ndarray:
    vector = np.frombuffer(
        blob,
        dtype=np.float32,
    )

    if (
        vector.shape[0]
        != dimensions
    ):
        raise RuntimeError(
            "Neplatný rozmer "
            "uloženého embeddingu"
        )

    return vector


def cosine_similarity(
    first: np.ndarray,
    second: np.ndarray,
) -> float:
    if (
        first.shape
        != second.shape
    ):
        raise ValueError(
            "Embeddingy majú "
            "rozdielny rozmer"
        )

    first_norm = float(
        np.linalg.norm(
            first
        )
    )

    second_norm = float(
        np.linalg.norm(
            second
        )
    )

    if (
        first_norm == 0.0
        or second_norm == 0.0
    ):
        return 0.0

    return float(
        np.dot(
            first,
            second,
        )
        / (
            first_norm
            * second_norm
        )
    )