Spaces:
Running
Running
File size: 3,019 Bytes
cfbbab4 | 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 | #!/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()
|