File size: 7,991 Bytes
ea26409 742ee3c ea26409 742ee3c ea26409 742ee3c ea26409 742ee3c ea26409 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """Render the model-card charts from results_full.json and internal_results.json into ../assets/."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt # noqa: E402
SCRIPT_DIR = Path(__file__).resolve().parent
ASSETS = SCRIPT_DIR.parent / "assets"
ORDER = ("mizan-v3", "mizan-v2", "gte-base", "bge-v2-m3")
LABELS = {
"mizan-v3": "Mizan-Rerank-v3 (306M)",
"mizan-v2": "Mizan-Rerank-V2 (306M)",
"gte-base": "gte-multilingual-reranker-base (306M)",
"bge-v2-m3": "bge-reranker-v2-m3 (568M)",
"retriever": "Retriever only (no reranker)",
}
COLORS = {"mizan-v3": "#0F766E", "mizan-v2": "#94A3B8", "gte-base": "#CBD5E1", "bge-v2-m3": "#F59E0B", "retriever": "#E2E8F0"}
DATASET_LABELS = {
"namaa_mrtydi": "NamaaMrTydi\n(MTEB, full)",
"namaa_mrtydi_unseen": "NamaaMrTydi\n(unseen subset)",
"arabic_hard_negatives": "Arabic Hard Negatives\n(full)",
"arabic_hard_negatives_unseen": "Arabic Hard Negatives\n(unseen subset)",
"gemini_test": "Adversarial short\n(Gemini test)",
"longctx_test": "Adversarial long-context\n(test)",
"multi_llm": "Multi-LLM\n(6 generators)",
}
HELD_OUT = ("namaa_mrtydi_unseen", "arabic_hard_negatives_unseen", "gemini_test", "longctx_test", "multi_llm")
def style(ax, title: str, ylabel: str) -> None:
ax.set_title(title, fontsize=13, fontweight="bold", loc="left")
ax.set_ylabel(ylabel)
ax.spines[["top", "right"]].set_visible(False)
ax.grid(axis="y", alpha=0.3)
ax.set_axisbelow(True)
def efficiency_chart(results: dict, speed: dict, path: Path) -> None:
quality = {key: sum(results[key]["datasets"][dataset]["ndcg@10"] for dataset in HELD_OUT) / len(HELD_OUT) for key in ORDER}
fig, ax = plt.subplots(figsize=(8.5, 4.6), dpi=150)
for key in ORDER:
row = speed["results"][key]
ax.scatter(row["pairs_per_second"], quality[key], s=row["parameters"] / 1e6 * 2.2, color=COLORS[key], edgecolor="#334155", linewidth=0.8, zorder=3)
offset = (-12, 0.006) if key == "gte-base" else (12, -0.004)
ax.annotate(f"{LABELS[key]}\n{quality[key]:.3f} nDCG@10 · {row['pairs_per_second']:.0f} pairs/s", (row["pairs_per_second"], quality[key]),
xytext=offset, textcoords="offset points", fontsize=8.5, ha="right" if key == "gte-base" else "left", va="center",
fontweight="bold" if key == "mizan-v3" else None)
rates = [speed["results"][key]["pairs_per_second"] for key in ORDER]
ax.set_xlim(min(rates) * 0.6, max(rates) * 1.5)
ax.set_ylim(min(quality.values()) - 0.03, max(quality.values()) + 0.03)
settings = speed["settings"]
ax.set_xlabel(f"throughput, query-passage pairs/s ({settings['device']}, {'fp16' if settings['fp16'] else 'fp32'}, max_length {settings['max_length']})")
style(ax, "Accuracy vs. speed (bubble size = parameters; up and right is better)", "held-out mean nDCG@10")
ax.grid(axis="x", alpha=0.3)
fig.tight_layout()
fig.savefig(path)
plt.close(fig)
def grouped_bars(results: dict, datasets: tuple[str, ...], metric: str, title: str, path: Path, floor: float) -> None:
fig, ax = plt.subplots(figsize=(11, 5.2), dpi=150)
width = 0.8 / len(ORDER)
for position, key in enumerate(ORDER):
values = [results[key]["datasets"][dataset][metric] for dataset in datasets]
bars = ax.bar([index + (position - (len(ORDER) - 1) / 2) * width for index in range(len(datasets))], values, width, label=LABELS[key], color=COLORS[key], edgecolor="white")
for bar, value in zip(bars, values):
ax.text(bar.get_x() + bar.get_width() / 2, value + 0.004, f"{value:.3f}", ha="center", va="bottom", fontsize=7.5, rotation=90)
ax.set_xticks(range(len(datasets)), [DATASET_LABELS[dataset] for dataset in datasets], fontsize=9)
ax.set_ylim(floor, 1.0)
style(ax, title, metric.replace("ndcg", "nDCG").replace("hit", "Hit").replace("mrr", "MRR"))
ax.legend(frameon=False, fontsize=8.5, ncol=2, loc="upper right")
fig.tight_layout()
fig.savefig(path)
plt.close(fig)
def summary_chart(results: dict, fiqh: dict, path: Path) -> None:
held_out = {key: sum(results[key]["datasets"][dataset]["ndcg@10"] for dataset in HELD_OUT) / len(HELD_OUT) for key in ORDER}
panels = (
("Held-out benchmarks\n(mean nDCG@10 over 5 sets)", held_out),
("Real fiqh search queries\n(nDCG@10, 224 queries)", {key: fiqh[key]["ndcg@10"] for key in ORDER}),
)
fig, axes = plt.subplots(1, 2, figsize=(12, 3.8), dpi=150)
for ax, (title, values) in zip(axes, panels):
keys = sorted(ORDER, key=lambda key: values[key])
bars = ax.barh([LABELS[key] for key in keys], [values[key] for key in keys], color=[COLORS[key] for key in keys])
for bar, key in zip(bars, keys):
ax.text(bar.get_width() + 0.003, bar.get_y() + bar.get_height() / 2, f"{values[key]:.3f}", va="center", fontsize=9, fontweight="bold" if key == "mizan-v3" else None)
ax.set_xlim(max(0.0, min(values.values()) - 0.08), max(values.values()) + 0.04)
ax.set_title(title, fontsize=11, fontweight="bold", loc="left")
ax.spines[["top", "right"]].set_visible(False)
ax.grid(axis="x", alpha=0.3)
ax.set_axisbelow(True)
fig.tight_layout()
fig.savefig(path)
plt.close(fig)
def training_chart(curve: dict, path: Path) -> None:
fig, ax = plt.subplots(figsize=(8, 4.2), dpi=150)
ax.plot(curve["epoch"], curve["longctx"], marker="o", color="#0F766E", label="Long-context adversarial (validation)")
ax.plot(curve["epoch"], curve["gemini"], marker="s", color="#F59E0B", label="Short adversarial (validation)")
best = max(range(len(curve["epoch"])), key=lambda index: curve["longctx"][index] + curve["gemini"][index])
ax.axvline(curve["epoch"][best], color="#64748B", linestyle="--", linewidth=1)
ax.text(curve["epoch"][best] - 0.1, curve["longctx"][best] - 0.012, f"released checkpoint\n(epoch {curve['epoch'][best]})", fontsize=8.5, color="#334155", ha="right", va="top")
ax.set_xticks(curve["epoch"], ["gte base"] + [str(epoch) for epoch in curve["epoch"][1:]])
ax.set_xlabel("epoch")
style(ax, "Validation nDCG@10 during training", "nDCG@10")
ax.legend(frameon=False, fontsize=9, loc="lower right")
fig.tight_layout()
fig.savefig(path)
plt.close(fig)
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--results", type=Path, default=SCRIPT_DIR / "results_full.json")
parser.add_argument("--internal", type=Path, default=SCRIPT_DIR / "internal_results.json")
parser.add_argument("--output-dir", type=Path, default=ASSETS)
args = parser.parse_args()
results = json.loads(args.results.read_text(encoding="utf-8"))["results"]
internal = json.loads(args.internal.read_text(encoding="utf-8"))
args.output_dir.mkdir(parents=True, exist_ok=True)
summary_chart(results, internal["fiqh_eval"]["models"], args.output_dir / "summary.png")
speed_file = SCRIPT_DIR / "speed_results.json"
if speed_file.is_file():
efficiency_chart(results, json.loads(speed_file.read_text(encoding="utf-8")), args.output_dir / "efficiency.png")
grouped_bars(results, ("namaa_mrtydi", "namaa_mrtydi_unseen", "arabic_hard_negatives_unseen"), "ndcg@10", "Public Arabic reranking benchmarks", args.output_dir / "public_benchmarks.png", 0.6)
grouped_bars(results, ("gemini_test", "longctx_test", "multi_llm"), "ndcg@10", "Adversarial Arabic reranking (held-out test sets)", args.output_dir / "adversarial_benchmarks.png", 0.5)
grouped_bars(results, HELD_OUT, "hit@1", "Hit@1: correct passage ranked first", args.output_dir / "hit_at_1.png", 0.0)
training_chart(internal["training_curve"], args.output_dir / "training_curve.png")
for path in sorted(args.output_dir.glob("*.png")):
print(path)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|