CLERC_FILE = (
"https://huggingface.co/datasets/jhu-clsp/CLERC/resolve/main/"
"teva_train_dir/train_data.jsonl.gz"
)
def cid(text: str) -> str:
"""Corpus id: a content hash, so passages shared across queries dedupe."""
return hashlib.sha1(text.encode("utf-8")).hexdigest()[:16]
@json_cache
def build_slice(n_rows: int, n_queries: int, seed: int) -> dict:
"""Stream CLERC rows, pool ``n_rows`` of them into a corpus, pick ``n_queries`` to evaluate."""
from datasets import load_dataset # heavy import, keep local
stream = load_dataset("json", data_files=CLERC_FILE, streaming=True, split="train")
rows = []
for row in stream:
if (
row.get("positive_passages")
and len(row.get("negative_passages") or []) == 20
):
rows.append(row)
if len(rows) >= 1000:
break
rng = random.Random(seed)
picked = rng.sample(rows, n_rows)
corpus, pool = {}, []
for row in picked:
gold = row["positive_passages"][0]["text"]
corpus[cid(gold)] = gold
for neg in row["negative_passages"]:
corpus[cid(neg["text"])] = neg["text"]
pool.append(
{"qid": str(row["query_id"]), "query": row["query"], "gold": cid(gold)}
)
# hold out the first 20 pooled rows; evaluate on the rest
queries = rng.sample(pool[20:], n_queries)
# sort the corpus by id so every run — live or cache replay — iterates it identically
return {"queries": queries, "corpus": dict(sorted(corpus.items()))}
def bm25_rankings(corpus: dict[str, str], queries: dict[str, str], k: int = 100):
"""Rank every passage in the corpus by word overlap with each query."""
import bm25s
cids = list(corpus)
retriever = bm25s.BM25()
retriever.index(bm25s.tokenize([corpus[c] for c in cids], stopwords="en"))
qids = list(queries)
idxs, _ = retriever.retrieve(
bm25s.tokenize([queries[q] for q in qids], stopwords="en"), k=min(k, len(cids))
)
return {q: [cids[i] for i in idxs[row]] for row, q in enumerate(qids)}
def gold_rank(ranked: list[str], gold: str) -> int | None:
"""1-based rank of the gold id, or None if it isn't in the list."""
return ranked.index(gold) + 1 if gold in ranked else None
SURFACE, INK, INK2, MUTED = "#f8f8f2", "#34342f", "#34342f", "#7c7c77"
GRID, AXIS, BLUE, GREEN = "#d8d8cf", "#d8d8cf", "#5d76a2", "#6f9b52"
def bar_chart(labels: list[str], shares: list[float], title: str) -> None:
"""A small single-series bar chart of shares (0-1, shown as percentages)."""
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(5, 3.2), facecolor=SURFACE)
ax.set_facecolor(SURFACE)
for side in ("top", "right"):
ax.spines[side].set_visible(False)
for side in ("left", "bottom"):
ax.spines[side].set_color(AXIS)
ax.tick_params(colors=MUTED, labelcolor=INK2, labelsize=9)
ax.set_axisbelow(True)
ax.grid(axis="y", color=GRID, linewidth=0.8)
bars = ax.bar(labels, shares, width=0.55, color=[BLUE, GREEN][: len(labels)])
ax.bar_label(
bars,
labels=[f"{s * 100:.0f}%" for s in shares],
padding=4,
color=INK,
fontsize=11,
)
ax.set_ylim(0, 1.1)
ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
ax.set_yticklabels(["0%", "25%", "50%", "75%", "100%"])
ax.set_ylabel(f"share of {len(queries)} queries", color=INK2, fontsize=9)
ax.set_title(title, loc="left", color=INK, fontsize=11)
plt.tight_layout()
display(fig)
plt.close(fig)
ds = build_slice(N_ROWS, N_QUERIES, seed=0)
corpus: dict[str, str] = ds["corpus"]
queries = {q["qid"]: q["query"] for q in ds["queries"]}
golds = {q["qid"]: q["gold"] for q in ds["queries"]}
candidates = {q: ranked[:TOP_K] for q, ranked in bm25_rankings(corpus, queries).items()}
in_top_k = sum(golds[q] in candidates[q] for q in queries)
at_rank_1 = sum(candidates[q][0] == golds[q] for q in queries)
bar_chart(
[f"In top {TOP_K}", "At rank 1"],
[in_top_k / len(queries), at_rank_1 / len(queries)],
f"Where the correct passage lands, {len(queries)} queries against {len(corpus):,} candidates",
)