Spaces:
Sleeping
Sleeping
| """ | |
| scripts/experiment.py | |
| --------------------- | |
| Experiment: k-sensitivity analysis — how does retrieval depth (k) affect | |
| Precision@k across all three retrieval strategies? | |
| This is a retrieval-only experiment — no OpenAI API calls are made. | |
| Results show how many correct chunks each retriever finds as k increases. | |
| Run (from project root): | |
| .venv/bin/python scripts/experiment.py catan | |
| .venv/bin/python scripts/experiment.py monopoly | |
| .venv/bin/python scripts/experiment.py all | |
| Outputs: | |
| data/outputs/{game}_experiment_k_sensitivity.csv | |
| """ | |
| import csv | |
| import json | |
| import os | |
| import sys | |
| import random | |
| from pathlib import Path | |
| K_VALUES = [1, 2, 3, 5, 7, 10] | |
| RETRIEVERS = ["random", "tfidf", "dense"] | |
| RANDOM_SEED = 42 | |
| # --------------------------------------------------------------------------- | |
| # Retrieval-only implementations (no OpenAI needed) | |
| # --------------------------------------------------------------------------- | |
| def _load_game(game: str) -> dict: | |
| """Load all retrieval artifacts for the given game.""" | |
| chunks_path = f"models/{game}/chunks.json" | |
| index_path = f"models/{game}/faiss.index" | |
| if not os.path.exists(chunks_path) or not os.path.exists(index_path): | |
| print(f"ERROR: Run build_features.py {game} first.", file=sys.stderr) | |
| sys.exit(1) | |
| with open(chunks_path) as f: | |
| chunks = json.load(f) | |
| try: | |
| import faiss | |
| from sentence_transformers import SentenceTransformer | |
| except ImportError as e: | |
| print(f"Missing dependency: {e}", file=sys.stderr) | |
| sys.exit(1) | |
| print(f" Loading dense model + FAISS index for '{game}'...") | |
| dense_model = SentenceTransformer("all-MiniLM-L6-v2") | |
| index = faiss.read_index(index_path) | |
| entry = {"chunks": chunks, "dense_model": dense_model, "index": index} | |
| tfidf_matrix_path = f"models/{game}/tfidf_matrix.npz" | |
| tfidf_vectorizer_path = f"models/{game}/tfidf_vectorizer.pkl" | |
| if os.path.exists(tfidf_matrix_path) and os.path.exists(tfidf_vectorizer_path): | |
| import pickle | |
| import scipy.sparse | |
| entry["tfidf_matrix"] = scipy.sparse.load_npz(tfidf_matrix_path) | |
| with open(tfidf_vectorizer_path, "rb") as f: | |
| entry["tfidf_vectorizer"] = pickle.load(f) | |
| print(f" Loaded TF-IDF index for '{game}'.") | |
| else: | |
| print(f" WARNING: TF-IDF index not found for '{game}' — skipping tfidf retriever.", file=sys.stderr) | |
| return entry | |
| def _retrieve_indices(query: str, k: int, retriever: str, reg: dict) -> list[int]: | |
| """Return the chunk indices selected by the given retriever for this query.""" | |
| chunks = reg["chunks"] | |
| if retriever == "random": | |
| rng = random.Random(RANDOM_SEED) | |
| return rng.sample(range(len(chunks)), min(k, len(chunks))) | |
| elif retriever == "tfidf": | |
| if "tfidf_vectorizer" not in reg: | |
| return [] | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| query_vec = reg["tfidf_vectorizer"].transform([query]) | |
| scores = cosine_similarity(query_vec, reg["tfidf_matrix"])[0] | |
| return list(scores.argsort()[::-1][:k]) | |
| else: # dense | |
| import faiss | |
| query_emb = reg["dense_model"].encode([query], convert_to_numpy=True) | |
| faiss.normalize_L2(query_emb) | |
| _, indices = reg["index"].search(query_emb, k) | |
| return [int(i) for i in indices[0] if i >= 0] | |
| def _precision_at_k(indices: list[int], keywords: list[str], chunks: list[dict]) -> float: | |
| """1 if at least one keyword appears in any retrieved chunk title, else 0.""" | |
| if not keywords: | |
| return 0.0 | |
| retrieved_titles = [chunks[i]["title"].lower() for i in indices] | |
| for kw in keywords: | |
| if any(kw.lower() in t for t in retrieved_titles): | |
| return 1.0 | |
| return 0.0 | |
| def _coverage(indices: list[int], keywords: list[str], chunks: list[dict]) -> float: | |
| """Fraction of required keywords found across retrieved chunks.""" | |
| if not keywords: | |
| return 0.0 | |
| retrieved_titles = [chunks[i]["title"].lower() for i in indices] | |
| hits = sum(1 for kw in keywords if any(kw.lower() in t for t in retrieved_titles)) | |
| return hits / len(keywords) | |
| # --------------------------------------------------------------------------- | |
| # Experiment runner | |
| # --------------------------------------------------------------------------- | |
| def run_experiment(game: str) -> list[dict]: | |
| eval_path = f"data/raw/{game}_eval.json" | |
| if not os.path.exists(eval_path): | |
| print(f"ERROR: {eval_path} not found.", file=sys.stderr) | |
| sys.exit(1) | |
| print(f"\nLoading artifacts for '{game}'...") | |
| reg = _load_game(game) | |
| with open(eval_path) as f: | |
| data = json.load(f) | |
| questions = data.get("correctness_questions", []) + data.get("stress_test_questions", []) | |
| print(f"\nRunning k-sensitivity experiment — {len(questions)} questions × {len(K_VALUES)} k-values × {len(RETRIEVERS)} retrievers") | |
| print("=" * 70) | |
| rows = [] | |
| for retriever in RETRIEVERS: | |
| if retriever == "tfidf" and "tfidf_vectorizer" not in reg: | |
| print(f" Skipping '{retriever}' — TF-IDF index not built.") | |
| continue | |
| for k in K_VALUES: | |
| precisions = [] | |
| coverages = [] | |
| for q in questions: | |
| indices = _retrieve_indices(q["question"], k, retriever, reg) | |
| keywords = q.get("source_keywords", []) | |
| precisions.append(_precision_at_k(indices, keywords, reg["chunks"])) | |
| coverages.append(_coverage(indices, keywords, reg["chunks"])) | |
| avg_prec = sum(precisions) / len(precisions) * 100 | |
| avg_cov = sum(coverages) / len(coverages) * 100 | |
| rows.append({ | |
| "game": game, | |
| "retriever": retriever, | |
| "k": k, | |
| "precision_at_k": round(avg_prec, 1), | |
| "coverage": round(avg_cov, 1), | |
| "n_questions": len(questions), | |
| }) | |
| print(f" [{retriever:<6}] k={k:2d} Precision@k={avg_prec:5.1f}% Coverage={avg_cov:5.1f}%") | |
| # Save CSV | |
| Path("data/outputs").mkdir(parents=True, exist_ok=True) | |
| out_path = f"data/outputs/{game}_experiment_k_sensitivity.csv" | |
| with open(out_path, "w", newline="") as f: | |
| writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) | |
| writer.writeheader() | |
| writer.writerows(rows) | |
| print(f"\nSaved → '{out_path}'") | |
| _print_experiment_summary(rows, game) | |
| return rows | |
| def _print_experiment_summary(rows: list[dict], game: str): | |
| print(f"\n{'='*70}") | |
| print(f"K-SENSITIVITY SUMMARY — {game.upper()}") | |
| print(f"{'='*70}") | |
| print(f"\n{'Retriever':<10} {'k':>4} {'Precision@k':>12} {'Coverage':>10}") | |
| print("-" * 44) | |
| for row in rows: | |
| print( | |
| f" {row['retriever']:<8} {row['k']:>4} " | |
| f"{row['precision_at_k']:>10.1f}% {row['coverage']:>8.1f}%" | |
| ) | |
| print() | |
| # Highlight key finding | |
| dense_k3 = next((r for r in rows if r["retriever"] == "dense" and r["k"] == 3), None) | |
| tfidf_k3 = next((r for r in rows if r["retriever"] == "tfidf" and r["k"] == 3), None) | |
| random_k3 = next((r for r in rows if r["retriever"] == "random" and r["k"] == 3), None) | |
| print("Key comparison at k=3:") | |
| for label, row in [("Dense (DL)", dense_k3), ("TF-IDF (ML)", tfidf_k3), ("Random", random_k3)]: | |
| if row: | |
| print(f" {label:<14} Precision@3={row['precision_at_k']:.1f}% Coverage={row['coverage']:.1f}%") | |
| # --------------------------------------------------------------------------- | |
| # Entry point | |
| # --------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser(description="k-sensitivity experiment for retriever comparison.") | |
| parser.add_argument("game", help="Game to run experiment on, or 'all' for both") | |
| args = parser.parse_args() | |
| games = ["catan", "monopoly"] if args.game == "all" else [args.game] | |
| for game in games: | |
| run_experiment(game) | |