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) )