Spaces:
Sleeping
Sleeping
File size: 3,018 Bytes
dfdc75f | 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 | """Download and sample BeIR/nfcorpus into local text pools for benchmarking.
nfcorpus is a small biomedical IR dataset whose two subsets give us the length contrast
SentenceTransformers warns about: short **queries** (latency workload) and long mixed-length
**corpus** documents (throughput workload).
Usage:
uv sync --extra bench
uv run python scripts/prepare_dataset.py # default n=1000, seed=42
Writes (gitignored):
benchmarks/data/nfcorpus-queries.jsonl short query text -> latency workload
benchmarks/data/nfcorpus-corpus.jsonl title + abstract -> throughput workload
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from datasets import load_dataset
from utils import length_stats, sample_texts
REPO_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = REPO_ROOT / "benchmarks" / "data"
MAX_CHARS = 8192 # mirror app MAX_CHARS_PER_TEXT so HTTP runs don't hit 422 on long docs
def build_pool(rows, to_text) -> list[str]:
"""Map dataset rows to cleaned, de-duplicated, length-capped strings."""
seen: set[str] = set()
pool: list[str] = []
for row in rows:
text = " ".join(to_text(row).split()) # collapse newlines/whitespace
if not text:
continue
text = text[:MAX_CHARS]
if text in seen:
continue
seen.add(text)
pool.append(text)
return pool
def write_jsonl(texts: list[str], path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as handle:
for text in texts:
handle.write(json.dumps({"text": text}) + "\n")
def main() -> None:
parser = argparse.ArgumentParser(description="Sample BeIR/nfcorpus into benchmark text pools.")
parser.add_argument("--n", type=int, default=1000, help="Max texts to keep per pool.")
parser.add_argument("--seed", type=int, default=42, help="Sampling seed (reproducible).")
args = parser.parse_args()
# queries -> short query text (latency workload)
queries = load_dataset("BeIR/nfcorpus", "queries")["queries"]
query_pool = build_pool(queries, lambda r: r["text"] or r["title"])
# corpus -> title + abstract (throughput workload, mixed/long)
corpus = load_dataset("BeIR/nfcorpus", "corpus")["corpus"]
corpus_pool = build_pool(corpus, lambda r: f"{r['title']}. {r['text']}")
query_sample = sample_texts(query_pool, min(args.n, len(query_pool)), args.seed)
corpus_sample = sample_texts(corpus_pool, min(args.n, len(corpus_pool)), args.seed)
write_jsonl(query_sample, DATA_DIR / "nfcorpus-queries.jsonl")
write_jsonl(corpus_sample, DATA_DIR / "nfcorpus-corpus.jsonl")
print(f"queries ({len(query_pool)} available) -> {length_stats(query_sample)}")
print(f"corpus ({len(corpus_pool)} available) -> {length_stats(corpus_sample)}")
print(f"Wrote {DATA_DIR}/nfcorpus-queries.jsonl and nfcorpus-corpus.jsonl")
if __name__ == "__main__":
main()
|