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