347 lines
7.5 KiB
Python
347 lines
7.5 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.text_budget import split_to_fit
|
|
|
|
|
|
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,
|
|
)
|
|
|
|
return SentenceTransformer(
|
|
model_name,
|
|
revision=revision,
|
|
)
|
|
|
|
|
|
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
|
|
)
|
|
)
|