1_embeddings.py

#!/usr/bin/env python3
"""Generate document embeddings. Full and reduced to 2d."""
import numpy as np
from sentence_transformers import SentenceTransformer
from umap import UMAP

from util import get_docs, EMBEDDINGS_PATH, config

embed_translation = config.get("EMBED_TRANSLATION", False)
docs = get_docs(translation=embed_translation in [True, "True", "true", "1", "on"])

sentence_model = SentenceTransformer(config.get("EMBEDDING_MODEL", "all-MiniLM-L6-v2"))
embeddings = sentence_model.encode(docs, show_progress_bar=True)
# save embeddings
with open(EMBEDDINGS_PATH, "wb") as f:
    np.save(f, embeddings, allow_pickle=False)


# save 2d embeddings
umap_model = UMAP(
    n_neighbors=15,
    n_components=2,
    min_dist=0.0,
    metric="cosine",
    low_memory=True,
    random_state=42,
)
doc_embeddings_reduced = umap_model.fit_transform(embeddings)
# save reduced embeddings
with open(EMBEDDINGS_PATH.with_suffix(".2d.npy"), "wb") as f:
    np.save(f, doc_embeddings_reduced, allow_pickle=False)