evaluate_threshold.py

import json
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# Embedding model used both for offline threshold evaluation and runtime cache
# lookup so that similarity scores are calibrated consistently.
MODEL_NAME = "all-MiniLM-L6-v2"

# Labeled evaluation set of query pairs.
DATASET_PATH = "query_pairs.json"

# Weighted error cost used for threshold selection.
# False positives are penalized more heavily because incorrectly reusing a
# cached answer is worse than making an extra LLM call.
FALSE_POSITIVE_COST = 10
FALSE_NEGATIVE_COST = 1

# Range of thresholds to sweep during evaluation.
THRESHOLD_START = 0.35
THRESHOLD_END = 0.95
THRESHOLD_STEP = 0.01


def load_pairs(path):
    """
    Load labeled query pairs from disk.

    Each pair is expected to include:
    - query_a
    - query_b
    - label
    """
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def get_similarity(model, text_a, text_b):
    """
    Compute cosine similarity between two query embeddings.
    """
    embeddings = model.encode([text_a, text_b], normalize_embeddings=True)
    sim = cosine_similarity([embeddings[0]], [embeddings[1]])[0][0]
    return float(sim)


def evaluate_threshold(pairs, threshold):
    """
    Evaluate a single similarity threshold on the labeled dataset.

    A pair is predicted as a cache hit when its similarity score is greater
    than or equal to the threshold.
    """
    true_positive = 0
    false_positive = 0
    true_negative = 0
    false_negative = 0

    for pair in pairs:
        score = pair["similarity"]
        actual_hit = pair["label"] == "should_hit"
        predicted_hit = score >= threshold

        if predicted_hit and actual_hit:
            true_positive += 1
        elif predicted_hit and not actual_hit:
            false_positive += 1
        elif not predicted_hit and not actual_hit:
            true_negative += 1
        else:
            false_negative += 1

    precision = 0.0
    recall = 0.0
    f1 = 0.0
    accuracy = 0.0

    if true_positive + false_positive > 0:
        precision = true_positive / (true_positive + false_positive)

    if true_positive + false_negative > 0:
        recall = true_positive / (true_positive + false_negative)

    if precision + recall > 0:
        f1 = 2 * precision * recall / (precision + recall)

    total = true_positive + false_positive + true_negative + false_negative
    if total > 0:
        accuracy = (true_positive + true_negative) / total

    false_positive_rate = 0.0
    if false_positive + true_negative > 0:
        false_positive_rate = false_positive / (false_positive + true_negative)

    cost = (
        FALSE_POSITIVE_COST * false_positive +
        FALSE_NEGATIVE_COST * false_negative
    )

    return {
        "threshold": threshold,
        "tp": true_positive,
        "fp": false_positive,
        "tn": true_negative,
        "fn": false_negative,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "accuracy": accuracy,
        "false_positive_rate": false_positive_rate,
        "cost": cost
    }


def print_top_mistakes(pairs, threshold, max_examples=10):
    """
    Print the most informative false positives and false negatives for the
    chosen threshold.

    False positives are sorted from highest similarity downward because these
    are the most dangerous incorrect cache hits.
    False negatives are sorted from lowest similarity upward because they are
    the strongest missed reuse opportunities near the decision boundary.
    """
    false_positives = []
    false_negatives = []

    for pair in pairs:
        score = pair["similarity"]
        actual_hit = pair["label"] == "should_hit"
        predicted_hit = score >= threshold

        if predicted_hit and not actual_hit:
            false_positives.append(pair)

        if not predicted_hit and actual_hit:
            false_negatives.append(pair)

    false_positives.sort(key=lambda pair: pair["similarity"], reverse=True)
    false_negatives.sort(key=lambda pair: pair["similarity"])

    print("\n" + "=" * 80)
    print("TOP FALSE POSITIVES")
    print("=" * 80)
    for pair in false_positives[:max_examples]:
        print("score:", round(pair["similarity"], 4))
        print("A:", pair["query_a"])
        print("B:", pair["query_b"])
        print("label:", pair["label"])
        print("-" * 80)

    print("\n" + "=" * 80)
    print("TOP FALSE NEGATIVES")
    print("=" * 80)
    for pair in false_negatives[:max_examples]:
        print("score:", round(pair["similarity"], 4))
        print("A:", pair["query_a"])
        print("B:", pair["query_b"])
        print("label:", pair["label"])
        print("-" * 80)


def main():
    """
    Run an offline threshold sweep over the labeled evaluation set.

    The script computes semantic similarity for each pair, evaluates a range
    of thresholds, and reports:
    - best threshold by weighted cost
    - best threshold by F1
    - top candidate thresholds by lowest cost
    - representative false positives and false negatives
    """
    print("Loading dataset...")
    pairs = load_pairs(DATASET_PATH)

    print("Loading embedding model...")
    model = SentenceTransformer(MODEL_NAME)

    print("Computing similarities...")
    for pair in pairs:
        pair["similarity"] = get_similarity(model, pair["query_a"], pair["query_b"])

    results = []
    threshold = THRESHOLD_START
    while threshold <= THRESHOLD_END:
        metrics = evaluate_threshold(pairs, round(threshold, 2))
        results.append(metrics)
        threshold += THRESHOLD_STEP

    best_by_cost = min(results, key=lambda row: row["cost"])
    best_by_f1 = max(results, key=lambda row: row["f1"])

    print("\nBEST THRESHOLD BY COST")
    print(best_by_cost)

    print("\nBEST THRESHOLD BY F1")
    print(best_by_f1)

    print("\nTOP 15 THRESHOLDS BY LOWEST COST")
    sorted_by_cost = sorted(results, key=lambda row: (row["cost"], -row["precision"]))
    for row in sorted_by_cost[:15]:
        print(
            "threshold=", row["threshold"],
            "cost=", row["cost"],
            "fp=", row["fp"],
            "fn=", row["fn"],
            "precision=", round(row["precision"], 4),
            "recall=", round(row["recall"], 4),
            "f1=", round(row["f1"], 4),
            "accuracy=", round(row["accuracy"], 4)
        )

    chosen_threshold = best_by_cost["threshold"]
    print_top_mistakes(pairs, chosen_threshold, max_examples=10)


if __name__ == "__main__":
    main()