Spaces:
Running on Zero
Running on Zero
| """Reproducible CPU/Arrow benchmark on one pinned Hugging Face shard. | |
| Install with ``python -m pip install ".[arrow,benchmark]"`` and accept the | |
| dataset's access terms before running this script. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import statistics | |
| import time | |
| from datetime import datetime, timezone | |
| import pyarrow as pa | |
| from datasets import load_dataset | |
| from huggingface_hub import HfApi, get_token, login | |
| from range_curate.examples.general_economy import DefaultTaxonomy, RelevanceEngine, __version__ | |
| DATASET = "Brianferrell787/financial-news-multisource" | |
| REVISION = "c25780f336280adb57c64bda7aed605d065c672d" | |
| SHARD = "data/sp500_daily_headlines/sp500_daily_headlines.000.parquet" | |
| def utc_timestamp(value: str) -> datetime: | |
| parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) | |
| if parsed.tzinfo is None: | |
| parsed = parsed.replace(tzinfo=timezone.utc) | |
| return parsed.astimezone(timezone.utc) | |
| def load_seed() -> tuple[pa.RecordBatch, str]: | |
| if get_token() is None: | |
| login() | |
| info = HfApi().dataset_info(DATASET, revision=REVISION, token=True) | |
| if info.sha != REVISION: | |
| raise RuntimeError(f"dataset revision drifted: {info.sha}") | |
| dataset = load_dataset( | |
| DATASET, | |
| name="data", | |
| data_files={"train": SHARD}, | |
| split="train", | |
| revision=REVISION, | |
| token=True, | |
| ).select_columns(["date", "text"]) | |
| texts = pa.array(dataset["text"], type=pa.string()) | |
| timestamps = pa.array( | |
| [utc_timestamp(value) for value in dataset["date"]], | |
| type=pa.timestamp("us", tz="UTC"), | |
| ) | |
| seed = pa.record_batch([texts, timestamps], names=["text", "timestamp"]) | |
| digest = hashlib.sha256() | |
| for date, text in zip(dataset["date"], dataset["text"], strict=True): | |
| row = json.dumps( | |
| [date, text], | |
| ensure_ascii=False, | |
| separators=(",", ":"), | |
| ).encode("utf-8") | |
| digest.update(len(row).to_bytes(8, "little")) | |
| digest.update(row) | |
| return seed, digest.hexdigest() | |
| def repeated_batches( | |
| seed: pa.RecordBatch, | |
| total_rows: int, | |
| *, | |
| batch_rows: int = 100_000, | |
| ) -> tuple[pa.RecordBatch, ...]: | |
| output = [] | |
| emitted = 0 | |
| while emitted < total_rows: | |
| start = emitted % seed.num_rows | |
| take = min(batch_rows, total_rows - emitted, seed.num_rows - start) | |
| output.append(seed.slice(start, take)) | |
| emitted += take | |
| return tuple(output) | |
| def reader(seed: pa.RecordBatch, batches: tuple[pa.RecordBatch, ...]): | |
| return pa.RecordBatchReader.from_batches(seed.schema, batches) | |
| def semantic_digest(facts, total_rows: int) -> str: | |
| digest = hashlib.sha256(f"rows={total_rows}\n".encode()) | |
| for row in facts.to_pyarrow().to_pylist(): | |
| payload = json.dumps( | |
| row, | |
| sort_keys=True, | |
| separators=(",", ":"), | |
| default=str, | |
| ).encode("utf-8") | |
| digest.update(len(payload).to_bytes(8, "little")) | |
| digest.update(payload) | |
| return digest.hexdigest() | |
| def measure( | |
| engine: RelevanceEngine, | |
| seed: pa.RecordBatch, | |
| total_rows: int, | |
| *, | |
| workers: int, | |
| repetitions: int, | |
| ) -> dict[str, object]: | |
| prepared = repeated_batches(seed, total_rows) | |
| warmup_rows = min(total_rows, 10_000) | |
| warmup = repeated_batches(seed, warmup_rows) | |
| engine.screen_arrow(reader(seed, warmup), workers=workers, row_count=warmup_rows) | |
| elapsed = [] | |
| digests = [] | |
| exceptions = [] | |
| for _ in range(repetitions): | |
| started = time.perf_counter() | |
| facts = engine.screen_arrow( | |
| reader(seed, prepared), | |
| workers=workers, | |
| row_count=total_rows, | |
| ) | |
| elapsed.append(time.perf_counter() - started) | |
| digests.append(semantic_digest(facts, total_rows)) | |
| exceptions.append(facts.exception_count) | |
| if len(set(digests)) != 1 or len(set(exceptions)) != 1: | |
| raise RuntimeError("repetitions produced different semantic outputs") | |
| median = statistics.median(elapsed) | |
| return { | |
| "backend": "cpu-rust-arrow", | |
| "rows": total_rows, | |
| "physical_rows": seed.num_rows, | |
| "replication_factor": round(total_rows / seed.num_rows, 3), | |
| "workers": workers, | |
| "repetitions": repetitions, | |
| "median_seconds": round(median, 6), | |
| "rows_per_second": round(total_rows / median), | |
| "exception_rows": exceptions[0], | |
| "output_sha256": digests[0], | |
| } | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| "--rows", | |
| type=int, | |
| nargs="+", | |
| default=[10_000, 100_000, 1_000_000], | |
| ) | |
| parser.add_argument("--workers", type=int, default=min(10, os.cpu_count() or 1)) | |
| parser.add_argument("--repetitions", type=int, default=3) | |
| args = parser.parse_args() | |
| seed, seed_sha256 = load_seed() | |
| taxonomy = DefaultTaxonomy() | |
| engine = RelevanceEngine(taxonomy=taxonomy, workers=args.workers) | |
| print( | |
| json.dumps( | |
| { | |
| "package": __version__, | |
| "taxonomy": taxonomy.fingerprint, | |
| "dataset": DATASET, | |
| "revision": REVISION, | |
| "shard": SHARD, | |
| "physical_rows": seed.num_rows, | |
| "seed_sha256": seed_sha256, | |
| }, | |
| indent=2, | |
| ) | |
| ) | |
| for total_rows in args.rows: | |
| print( | |
| json.dumps( | |
| measure( | |
| engine, | |
| seed, | |
| total_rows, | |
| workers=args.workers, | |
| repetitions=args.repetitions, | |
| ), | |
| sort_keys=True, | |
| ) | |
| ) | |
| if __name__ == "__main__": | |
| main() | |