"""
Modul pro vytváření sémantických embeddingů a dalších funkcí.
"""
import numpy as np
from sentence_transformers import SentenceTransformer
import nltk
nltk.download("punkt")
from nltk import tokenize
[docs]
def get_embeddings(text: str, model: SentenceTransformer, sent_tokenize: bool = True) -> np.ndarray:
"""Převede text na sémantické embeddingy (vektor).
Umožňuje převést kratší i delší text na embeddingy. Pokud je sent_tokenize True, tak nejdříve zákoduje jednotlivé \
věty a pak jednotlivé věty zprůměruje k zákodování celého textu. V případě, že je False, tak zákoduje text jako \
celek. Druhá varianta lépe zachytí celkový kontext, ale většina SentenceTransformer modelů umí najednou zpracovat \
pouze určité množství tokenů (např. 128 tokenů, což to odpovídá cca třem průměrným větám), po tomto množství \
kvalita klesá.
:param text: text, který chceme zakódvat do embeddingů
:param model: SentenceTransformer model, podle kterého se text zakóduje do embeddingů
:param sent_tokenize: pokud je nastavený, tak nejdříve zakóduje jednotlivé věty a až pak celý text.
:return: embeddingový vektor vstupního textu
"""
if sent_tokenize:
sentences = tokenize.sent_tokenize(text)
sentence_embeddings = [model.encode(sent) for sent in sentences]
doc_embedding = np.mean(sentence_embeddings, axis=0)
else:
doc_embedding = model.encode(text)
return doc_embedding