Spaces:
Running
Running
File size: 4,260 Bytes
a66a15c | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | #!/usr/bin/env python3
"""Build the FinQA viewer data from the raw dataset splits.
Reads the four raw split files (train/dev/test/private_test.json) and writes:
eval.jsonl one slim record per QA example (9,200 lines)
corpus/index.json light per-document index (2,789 unique pages)
corpus/records/<key>.json full page content (pre_text + table + post_text)
A "document" is a unique source page identified by `filename`
(`TICKER/YEAR/page_N.pdf`). Page content is identical across all examples that
share a filename, so it is de-duplicated into the corpus.
Run from the viewer repo root:
python scripts/build_data.py [--src <dataset-dir>]
Default --src is $DATA_ROOT/finqa/dataset (falls back to ./dataset).
"""
import argparse
import json
import os
from collections import Counter, defaultdict
SPLITS = ["train", "dev", "test", "private_test"]
# qa fields worth surfacing in the Eval tab (heavy retriever artifacts dropped).
QA_KEEP = [
"question", "answer", "exe_ans", "program", "program_re", "steps",
"gold_inds", "explanation", "ann_table_rows", "ann_text_rows",
]
def safe_key(filename: str) -> str:
"""`ADI/2009/page_49.pdf` -> `ADI__2009__page_49.pdf` (filesystem-safe)."""
return filename.replace("/", "__")
def main() -> None:
root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
default_src = os.path.join(
os.environ.get("DATA_ROOT", "/mnt/ramdisk/blobstore/timchen0618/data"),
"finqa", "dataset",
)
if not os.path.isdir(default_src):
default_src = os.path.join(root, "dataset")
ap = argparse.ArgumentParser()
ap.add_argument("--src", default=default_src, help="dir with the raw split json files")
args = ap.parse_args()
records_dir = os.path.join(root, "corpus", "records")
os.makedirs(records_dir, exist_ok=True)
# clear stale shards
for fn in os.listdir(records_dir):
if fn.endswith(".json"):
os.remove(os.path.join(records_dir, fn))
corpus: dict[str, dict] = {}
counts: dict[str, Counter] = defaultdict(Counter)
eval_rows: list[dict] = []
for split in SPLITS:
path = os.path.join(args.src, f"{split}.json")
with open(path) as f:
data = json.load(f)
for x in data:
fn = x["filename"]
counts[fn][split] += 1
if fn not in corpus:
ticker, year, page = fn.split("/")
corpus[fn] = {
"filename": fn,
"ticker": ticker,
"year": year,
"page": page.replace("page_", "").replace(".pdf", ""),
"pre_text": x["pre_text"],
"table": x["table"],
"post_text": x["post_text"],
}
qa = x.get("qa", {})
row = {"id": x["id"], "filename": fn, "split": split}
for k in QA_KEEP:
if k in qa:
row[k] = qa[k]
eval_rows.append(row)
# eval.jsonl
with open(os.path.join(root, "eval.jsonl"), "w") as f:
for row in eval_rows:
f.write(json.dumps(row, ensure_ascii=False) + "\n")
# corpus records + index
index = []
for fn, doc in corpus.items():
doc["num_questions"] = sum(counts[fn].values())
doc["splits"] = sorted(counts[fn].keys())
key = safe_key(fn)
with open(os.path.join(records_dir, key + ".json"), "w") as f:
json.dump(doc, f, ensure_ascii=False)
index.append({
"filename": fn,
"key": key,
"ticker": doc["ticker"],
"year": doc["year"],
"page": doc["page"],
"num_questions": doc["num_questions"],
"splits": doc["splits"],
})
def sort_key(d):
return (d["ticker"], d["year"], int(d["page"]) if d["page"].isdigit() else 0)
index.sort(key=sort_key)
with open(os.path.join(root, "corpus", "index.json"), "w") as f:
json.dump(index, f, ensure_ascii=False)
print(f"eval.jsonl: {len(eval_rows)} rows")
print(f"corpus: {len(index)} docs -> corpus/records/*.json + corpus/index.json")
if __name__ == "__main__":
main()
|