Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Build lazy compressed text shards for the FinanceBench Corpus tab.""" | |
| from __future__ import annotations | |
| import argparse | |
| import base64 | |
| import gzip | |
| import json | |
| import shutil | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from typing import Any, Iterator | |
| ROOT = Path(__file__).resolve().parent.parent | |
| DEFAULT_CORPUS = Path( | |
| "/mnt/ramdisk/blobstore/timchen0618/data/eval/financebench/" | |
| "unified/corpus.unified.jsonl" | |
| ) | |
| def iter_jsonl(path: Path) -> Iterator[dict[str, Any]]: | |
| with path.open(encoding="utf-8") as handle: | |
| for line in handle: | |
| if line.strip(): | |
| yield json.loads(line) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS) | |
| parser.add_argument("--out", type=Path, default=ROOT / "corpus_text") | |
| args = parser.parse_args() | |
| grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) | |
| metadata: dict[str, dict[str, Any]] = {} | |
| for row in iter_jsonl(args.corpus): | |
| meta = row.get("meta") or {} | |
| doc_name = str(meta["doc_name"]) | |
| metadata.setdefault(doc_name, dict(meta)) | |
| grouped[doc_name].append( | |
| { | |
| "id": str(row["id"]), | |
| "page_start": meta.get("page_start"), | |
| "page_end": meta.get("page_end"), | |
| "contents": str(row.get("contents") or ""), | |
| } | |
| ) | |
| records_dir = args.out / "records" | |
| if records_dir.exists(): | |
| shutil.rmtree(records_dir) | |
| records_dir.mkdir(parents=True, exist_ok=True) | |
| index_rows = [] | |
| for doc_name in sorted(grouped): | |
| chunks = sorted(grouped[doc_name], key=lambda row: row["page_start"] or 0) | |
| record = { | |
| "doc_name": doc_name, | |
| "meta": metadata[doc_name], | |
| "chunks": chunks, | |
| } | |
| raw = json.dumps(record, ensure_ascii=False, separators=(",", ":")).encode() | |
| encoded = base64.b64encode(gzip.compress(raw, compresslevel=9)).decode("ascii") | |
| path = f"records/{doc_name}.json.gzip.b64" | |
| (args.out / path).write_text(encoded, encoding="ascii") | |
| index_rows.append( | |
| { | |
| "doc_name": doc_name, | |
| "n_chunks": len(chunks), | |
| "n_pages": max(chunk["page_end"] or 0 for chunk in chunks), | |
| "path": path, | |
| } | |
| ) | |
| args.out.mkdir(parents=True, exist_ok=True) | |
| (args.out / "index.json").write_text( | |
| json.dumps( | |
| { | |
| "meta": { | |
| "n_documents": len(index_rows), | |
| "n_chunks": sum(row["n_chunks"] for row in index_rows), | |
| "source": str(args.corpus), | |
| }, | |
| "rows": index_rows, | |
| }, | |
| ensure_ascii=False, | |
| ), | |
| encoding="utf-8", | |
| ) | |
| print(f"Wrote {len(index_rows)} documents to {args.out}") | |
| if __name__ == "__main__": | |
| main() | |