scripts/build_data.py

#!/usr/bin/env python3
"""Build corpus_index.json for the FinanceBench viewer.

Combines the PDF listing (the actual corpus documents) with the per-document
metadata in financebench_document_information.jsonl. Every PDF becomes one
corpus entry; metadata is attached when available (8 PDFs have no metadata row).

Run from the viewer repo root:
    python scripts/build_data.py

Reads:
    financebench_document_information.jsonl   (doc metadata)
    pdfs/                                      (symlink to the PDF corpus)
Writes:
    corpus_index.json                          (sorted list of doc entries)
"""
import json
import os

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DOCINFO = os.path.join(ROOT, "financebench_document_information.jsonl")
PDFS = os.path.join(ROOT, "pdfs")
OUT = os.path.join(ROOT, "corpus_index.json")


def main() -> None:
    meta = {}
    with open(DOCINFO) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            meta[row["doc_name"]] = row

    pdf_names = sorted(
        fn[:-4] for fn in os.listdir(PDFS) if fn.lower().endswith(".pdf")
    )

    corpus = []
    for name in pdf_names:
        m = meta.get(name, {})
        corpus.append(
            {
                "doc_name": name,
                "company": m.get("company"),
                "gics_sector": m.get("gics_sector"),
                "doc_type": m.get("doc_type"),
                "doc_period": m.get("doc_period"),
                "doc_link": m.get("doc_link"),
                "has_meta": name in meta,
                "pdf": f"pdfs/{name}.pdf",
            }
        )

    with open(OUT, "w") as f:
        json.dump(corpus, f, ensure_ascii=False, indent=0)

    with_meta = sum(1 for c in corpus if c["has_meta"])
    print(f"wrote {OUT}: {len(corpus)} docs ({with_meta} with metadata)")


if __name__ == "__main__":
    main()