Spaces:
Running
Running
| """Plot the retrieval comparison per embedder from retrieval_results.json. | |
| Generates docs/retrieval-embedders.svg: overall MRR (bars grouped by architecture, | |
| one cluster per embedder) + Vector hit@k curves (embedder effect). | |
| Requires the [notebooks] extra (matplotlib). | |
| python -m eval.plot_retrieval | |
| """ | |
| import json | |
| from pathlib import Path | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| ROOT = Path(__file__).resolve().parent.parent | |
| COLORS = {"vector": "#3b82f6", "hybrid": "#22c55e", "graph": "#a855f7"} | |
| SHORT = {"vector": "Vector", "hybrid": "Hybrid", "graph": "Graph"} | |
| def _kind(name: str) -> str: | |
| n = name.lower() | |
| if "vector" in n or "vecto" in n: | |
| return "vector" | |
| if "hybr" in n: | |
| return "hybrid" | |
| return "graph" | |
| def main() -> None: | |
| data = json.loads((ROOT / "eval" / "retrieval_results.json").read_text("utf-8")) | |
| embedders = data["config"]["embedders"] | |
| ks = data["config"]["ks"] | |
| results = data["results"] | |
| stacks = list(results[embedders[0]]["stacks"]) | |
| kinds = [_kind(s) for s in stacks] | |
| labels = [e.split("/")[-1] for e in embedders] | |
| fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.5)) | |
| x = np.arange(len(embedders)) | |
| width = 0.25 | |
| for i, (s, k) in enumerate(zip(stacks, kinds)): | |
| vals = [results[e]["stacks"][s]["overall"]["mrr"] for e in embedders] | |
| bars = ax1.bar(x + (i - 1) * width, vals, width, label=SHORT[k], color=COLORS[k]) | |
| ax1.bar_label(bars, fmt="%.2f", fontsize=8, padding=2) | |
| ax1.set_xticks(x) | |
| ax1.set_xticklabels(labels, fontsize=9) | |
| ax1.set_ylabel("MRR (overall)") | |
| ax1.set_ylim(0, 1) | |
| ax1.set_title("MRR by embedder and architecture") | |
| ax1.legend(fontsize=8) | |
| ax1.grid(axis="y", alpha=0.3) | |
| vec = stacks[kinds.index("vector")] | |
| styles = ["-o", "--s", ":^"] | |
| for e, label, style in zip(embedders, labels, styles): | |
| ys = [results[e]["stacks"][vec]["overall"][f"hit@{k}"] for k in ks] | |
| ax2.plot(ks, ys, style, label=label) | |
| ax2.set_xlabel("k") | |
| ax2.set_ylabel("hit@k") | |
| ax2.set_ylim(0, 1.02) | |
| ax2.set_xticks(ks) | |
| ax2.set_title("Vector: hit@k by embedder") | |
| ax2.legend(fontsize=8) | |
| ax2.grid(alpha=0.3) | |
| n_q = data["config"].get("n_questions", "?") | |
| n_art = data["config"].get("n_articles", "?") | |
| fig.suptitle(f"The embedding model changes retrieval ({n_q} questions, {n_art} articles)", | |
| fontsize=12) | |
| fig.tight_layout() | |
| out = ROOT / "docs" / "retrieval-embedders.svg" | |
| fig.savefig(out, bbox_inches="tight") | |
| print(f"✅ {out}") | |
| if __name__ == "__main__": | |
| main() | |