from sentence_transformers import SentenceTransformer


model = SentenceTransformer("models/support-matryoshka/final")
texts = [
    "reset admin login",
    "help an administrator regain account access",
    "retrieve the newest billing invoice",
]

for dimension in [384, 128, 64, 32]:
    embeddings = model.encode(
        texts,
        normalize_embeddings=True,
        truncate_dim=dimension,
    )
    scores = embeddings[1:] @ embeddings[0]
    assert embeddings.shape == (3, dimension)
    assert scores[0] > scores[1]
    print(
        f"{dimension:>3} dimensions: "
        f"related={scores[0]:.4f} unrelated={scores[1]:.4f}"
    )