285 lines
7.1 KiB
Python
285 lines
7.1 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import hashlib
|
|
import json
|
|
from functools import lru_cache
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from tools.retrieval.text_budget import split_to_fit
|
|
from tools.runtime.runtime_assets import local_model
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
def get_embedding_model(
|
|
model_name: str,
|
|
):
|
|
return _load_embedding_model(model_name, os.getenv("EMBEDDING_REVISION") or None)
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def _load_embedding_model(model_name: str, revision: str | None):
|
|
from sentence_transformers import (
|
|
SentenceTransformer,
|
|
)
|
|
|
|
bundled = local_model(model_name, revision)
|
|
return SentenceTransformer(
|
|
bundled or model_name,
|
|
revision=None if bundled else revision,
|
|
local_files_only=bool(bundled) or os.getenv("HF_HUB_OFFLINE") == "1",
|
|
)
|
|
|
|
|
|
def model_profile(model_name: str) -> dict:
|
|
"""Fingerprint the loaded weights, tokenizer and pooling, not a mutable alias."""
|
|
return _model_profile(model_name, get_embedding_model(model_name))
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def _model_profile(model_name, model):
|
|
import torch
|
|
|
|
weights = hashlib.sha256()
|
|
for name, tensor in sorted(model.state_dict().items()):
|
|
weights.update(name.encode())
|
|
weights.update(str((tuple(tensor.shape), tensor.dtype)).encode())
|
|
weights.update(tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes())
|
|
tokenizer_config = json.loads(model.tokenizer.backend_tokenizer.to_str())
|
|
# Fast tokenizers store the last call's padding/truncation in the backend.
|
|
# These mutable execution settings must not change the vocabulary fingerprint.
|
|
tokenizer_config.pop("padding", None)
|
|
tokenizer_config.pop("truncation", None)
|
|
tokenizer = json.dumps(
|
|
tokenizer_config, sort_keys=True, separators=(",", ":"), ensure_ascii=False
|
|
)
|
|
modules = {
|
|
name: module.get_config_dict()
|
|
for name, module in model._modules.items()
|
|
if hasattr(module, "get_config_dict")
|
|
}
|
|
profile = {
|
|
"model": model_name,
|
|
"weights_sha256": weights.hexdigest(),
|
|
"tokenizer_sha256": hashlib.sha256(tokenizer.encode()).hexdigest(),
|
|
"max_seq_length": int(model.max_seq_length),
|
|
"modules": modules,
|
|
"passage_prefix": "passage: ",
|
|
"query_prefix": "query: ",
|
|
"normalize_embeddings": True,
|
|
"dtype": "float32",
|
|
"segment_policy": "lossless-v1",
|
|
}
|
|
profile["fingerprint"] = hashlib.sha256(
|
|
json.dumps(profile, sort_keys=True).encode()
|
|
).hexdigest()
|
|
return profile
|
|
|
|
|
|
def passage_token_count(text: str, *, model_name: str | None = None) -> int:
|
|
model = get_embedding_model(model_name or embedding_model_name())
|
|
return len(
|
|
model.tokenizer(
|
|
"passage: " + text, truncation=False, add_special_tokens=True, verbose=False
|
|
)["input_ids"]
|
|
)
|
|
|
|
|
|
def split_passage(text: str, *, model_name: str | None = None) -> list[str]:
|
|
selected = model_name or embedding_model_name()
|
|
limit = int(get_embedding_model(selected).max_seq_length)
|
|
return split_to_fit(
|
|
text, lambda value: passage_token_count(value, model_name=selected) <= limit
|
|
)
|
|
|
|
|
|
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 for text in texts]
|
|
if any(
|
|
passage_token_count(text, model_name=selected_model) > model.max_seq_length
|
|
for text in texts
|
|
):
|
|
raise ValueError("Embedding input exceeds model limit; split it before encoding.")
|
|
|
|
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)
|
|
|
|
length = len(
|
|
model.tokenizer(
|
|
"query: " + clean_query, truncation=False, add_special_tokens=True, verbose=False
|
|
)["input_ids"]
|
|
)
|
|
if length > model.max_seq_length:
|
|
raise ValueError("Query prekračuje tokenový limit embeddingového modelu")
|
|
|
|
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)
|
|
)
|