test_fastembed_threshold.py

from fastembed import TextEmbedding
import numpy as np

model = TextEmbedding()

def cos_sim(a, b):
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

pairs = [
    # ---- genuine paraphrases: SHOULD score high ----
    ("paraphrase", "what caused the checkout incident", "why did checkout break in June"),
    ("paraphrase", "what are the risks for Northwind's renewal", "what could jeopardize the Northwind deal"),
    ("paraphrase", "how do I reset my password", "what's the process to change my login credentials"),
    ("paraphrase", "when is the next deploy freeze", "what dates are we not allowed to ship"),
    ("paraphrase", "who owns the sharding project", "who's responsible for the index sharding rollout"),

    # ---- genuinely unrelated: SHOULD score low ----
    ("unrelated", "what caused the checkout incident", "when is the next deploy freeze"),
    ("unrelated", "how do I reset my password", "who owns the sharding project"),
    ("unrelated", "what are the risks for Northwind's renewal", "what's the weather like today"),
    ("unrelated", "when is the next deploy freeze", "how do I reset my password"),

    # ---- the hard case: same SHAPE, different SUBJECT — should NOT merge ----
    ("same-shape-diff-subject", "what caused the checkout incident", "what caused the Northwind attachment bug"),
    ("same-shape-diff-subject", "what are the risks for Northwind's renewal", "what are the risks for Acme's renewal"),
    ("same-shape-diff-subject", "how do I reset my password", "how do I reset my API token"),
    ("same-shape-diff-subject", "who owns the sharding project", "who owns the MFA rollout"),
]

results = []
for category, q1, q2 in pairs:
    e1, e2 = list(model.embed([q1, q2]))
    score = cos_sim(e1, e2)
    results.append((category, score, q1, q2))

results.sort(key=lambda r: -r[1])

print(f"{'score':>6}  {'category':<24}  pair")
print("-" * 90)
for category, score, q1, q2 in results:
    print(f"{score:.3f}   {category:<24}  \"{q1}\" <-> \"{q2}\"")

print()
by_cat = {}
for category, score, _, _ in results:
    by_cat.setdefault(category, []).append(score)
for cat, scores in by_cat.items():
    print(f"{cat:<24} min={min(scores):.3f}  max={max(scores):.3f}  avg={sum(scores)/len(scores):.3f}")