Create evaluate.py

evaluate.py

@@ -0,0 +1,76 @@
+"""
+evaluate.py
+
+Reproduce the retrieval metrics.
+"""
+
+import numpy as np
+from collections import Counter
+
+from hv_temporal_context import HVTemporalContext, DEFAULT_K
+from train_temporal_context import load_turns, _fill, TOPIC_TEMPLATES
+
+
+def generate_queries(seed=99, n_per_topic=4):
+    rng = np.random.default_rng(seed)
+    queries = []
+    for topic, tpl in TOPIC_TEMPLATES.items():
+        for _ in range(n_per_topic):
+            u = _fill(rng.choice(tpl["user"]), tpl["fills"], rng)
+            queries.append((topic, u))
+    rng.shuffle(queries)
+    return queries
+
+
+def evaluate_similarity(model, queries, k=DEFAULT_K):
+    precisions = []
+    for true_topic, query in queries:
+        idx, _ = model.retrieve(query, k)
+        topics = [model.topics[i] for i in idx]
+        precisions.append(sum(1 for t in topics if t == true_topic) / len(topics))
+    return float(np.mean(precisions))
+
+
+def evaluate_temporal(model, queries, k=DEFAULT_K, lambda_surprise=0.5):
+    precisions = []
+    for true_topic, query in queries:
+        idx, _ = model.retrieve_temporal(query, true_topic, k, lambda_surprise)
+        topics = [model.topics[i] for i in idx]
+        precisions.append(sum(1 for t in topics if t == true_topic) / len(topics))
+    return float(np.mean(precisions))
+
+
+def evaluate_random(model, queries, k=DEFAULT_K, seed=0):
+    rng = np.random.default_rng(seed)
+    precisions = []
+    for true_topic, _ in queries:
+        idx = rng.choice(len(model.topics), size=k, replace=False)
+        topics = [model.topics[i] for i in idx]
+        precisions.append(sum(1 for t in topics if t == true_topic) / len(topics))
+    return float(np.mean(precisions))
+
+
+def main(model_path="temporal_context_model"):
+    print("Loading model...")
+    model = HVTemporalContext.load(model_path)
+    print(f"  {len(model.texts)} turns, "
+          f"{len(model.vocab_index)} vocabulary")
+
+    queries = generate_queries(seed=99, n_per_topic=4)
+
+    p_rand = evaluate_random(model, queries)
+    p_sim = evaluate_similarity(model, queries)
+    p_temp = evaluate_temporal(model, queries)
+
+    print()
+    print(f"  {'method':<28} {'precision@5':>12}")
+    print("  " + "-" * 42)
+    print(f"  {'random':<28} {p_rand:>12.3f}")
+    print(f"  {'similarity':<28} {p_sim:>12.3f}")
+    print(f"  {'position-conditional':<28} {p_temp:>12.3f}")
+    print()
+    print(f"  gap over similarity:  {p_temp - p_sim:+.3f}")
+
+
+if __name__ == "__main__":
+    main()