from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka")
sentences = [
    "A developer shortens embeddings for a vector database.",
    "An engineer reduces vector dimensions for semantic search storage.",
    "The invoice was paid by bank transfer yesterday.",
]

full_embeddings = model.encode(sentences, normalize_embeddings=True)
truncated_embeddings = model.encode(
    sentences,
    normalize_embeddings=True,
    truncate_dim=128,
)

print(f"full embedding shape: {full_embeddings.shape}")
print(f"truncated embedding shape: {truncated_embeddings.shape}")
print(f"related pair cosine, full: {np.dot(full_embeddings[0], full_embeddings[1]):.4f}")
print(f"related pair cosine, truncated: {np.dot(truncated_embeddings[0], truncated_embeddings[1]):.4f}")
print(f"unrelated pair cosine, truncated: {np.dot(truncated_embeddings[0], truncated_embeddings[2]):.4f}")