scripts/build_data.py

#!/usr/bin/env python3
"""Build viewer data for the CorpusQA viewer.

CorpusQA ships one raw JSONL per context-length "set" (128k, 1m, ...). Each line
is a self-contained QA instance whose user prompt concatenates the supporting
documents (delimited by ``# Document N:`` markers) followed by a ``# Question:``
block. The same source documents are reused across many questions, so the true
corpus is small (e.g. 23 unique docs in the 128k set).

For each set this script derives two compact, browser-friendly files:

  corpus_<set>.json  list[{title, domain, ext, size, n_questions, content}]
  eval_<set>.json    list[{id, domain, type, question, answer, doc_files,
                            system_prompt, answer_format}]

and upserts a ``sets.json`` manifest the viewer reads to populate its set
selector. The raw JSONL is NOT copied into the Space (it is 108MB / 1GB).

Run from the viewer repo root:

    python scripts/build_data.py --set 128k \
        --input /mnt/ramdisk/blobstore/timchen0618/data/corpusqa/128k_4domains.jsonl
    python scripts/build_data.py --set 1m \
        --input /mnt/ramdisk/blobstore/timchen0618/data/corpusqa/1m_4domains.jsonl
"""
import argparse
import json
import os
import re
import shutil

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_DATA_DIR = "/mnt/ramdisk/blobstore/timchen0618/data/corpusqa"

# Splits the docs region into per-document chunks. The marker looks like
# "# Document 12:" optionally followed by whitespace/newline.
DOC_MARKER_RE = re.compile(r"# Document \d+:\s*")
# The question block starts at the first "# Question:" marker (case as emitted
# by the dataset). Everything before it is the concatenated documents.
QUESTION_SPLIT_RE = re.compile(r"# Question:\s*", re.IGNORECASE)


def user_content(prompt):
    """Return the concatenated user-message content of a record's prompt."""
    parts = [m.get("content", "") for m in prompt if m.get("role") == "user"]
    return "\n".join(parts)


def system_content(prompt):
    parts = [m.get("content", "") for m in prompt if m.get("role") == "system"]
    return "\n".join(parts).strip()


def split_documents(user):
    """Split a user prompt into (list_of_doc_texts, question_block).

    The question block is everything from the first ``# Question:`` marker
    onwards (with the marker stripped); the documents are the ``# Document N:``
    chunks that precede it.
    """
    m = QUESTION_SPLIT_RE.search(user)
    if m:
        docs_region = user[: m.start()]
        question_block = user[m.end():].strip()
    else:
        docs_region, question_block = user, ""
    chunks = DOC_MARKER_RE.split(docs_region)
    # chunks[0] is any preamble before "# Document 1:" (normally empty).
    docs = [c.strip() for c in chunks[1:]]
    return docs, question_block


def ext_of(title):
    base = title.rsplit(".", 1)
    return base[1].lower() if len(base) == 2 else ""


def slugify(idx, title):
    """Deterministic, collision-free ASCII shard filename for a doc title.

    Titles may be non-ASCII (Chinese) or contain spaces; the numeric prefix
    guarantees uniqueness even when the ASCII slug collapses to the same value.
    """
    safe = re.sub(r"[^0-9A-Za-z._-]+", "_", title).strip("_")[:80] or "doc"
    return f"{idx:04d}_{safe}"


def build_set(set_name, input_path):
    # corpus keyed by title -> {content (longest), domain, n_questions}
    corpus = {}
    eval_rows = []
    n_records = 0
    n_mismatch = 0
    n_varied = 0

    with open(input_path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            r = json.loads(line)
            n_records += 1
            doc_files = r.get("doc_files") or []
            user = user_content(r.get("prompt") or [])
            docs, question_block = split_documents(user)

            if len(docs) != len(doc_files):
                n_mismatch += 1
                # Fall back to whatever pairs up; still record the question.
            for title, content in zip(doc_files, docs):
                entry = corpus.get(title)
                if entry is None:
                    corpus[title] = {
                        "domain": r.get("domain", ""),
                        "content": content,
                        "n_questions": 1,
                    }
                else:
                    entry["n_questions"] += 1
                    if content != entry["content"]:
                        n_varied += 1
                        # Keep the longest occurrence as canonical.
                        if len(content) > len(entry["content"]):
                            entry["content"] = content

            # Derive the answer-format instructions: the question block minus
            # the (already-separate) question text.
            question = (r.get("question") or "").strip()
            answer_format = question_block
            if question and question_block.startswith(question):
                answer_format = question_block[len(question):].strip()

            eval_rows.append({
                "id": r.get("id", ""),
                "domain": r.get("domain", ""),
                "type": r.get("type", ""),
                "question": question,
                "answer": r.get("answer"),
                "doc_files": doc_files,
                "system_prompt": system_content(r.get("prompt") or []),
                "answer_format": answer_format,
            })

    corpus_rows = [
        {
            "title": title,
            "domain": e["domain"],
            "ext": ext_of(title),
            "size": len(e["content"]),
            "n_questions": e["n_questions"],
            "content": e["content"],
        }
        for title, e in corpus.items()
    ]
    corpus_rows.sort(key=lambda d: (d["domain"], d["title"]))
    eval_rows.sort(key=lambda d: (d["domain"], d["id"]))

    # Write per-document content shards (avoids a single >10MB corpus file and
    # lets the viewer lazy-load one document at a time). The index carries only
    # metadata + a pointer to each shard.
    shard_dir = os.path.join(ROOT, f"corpus_{set_name}")
    if os.path.isdir(shard_dir):
        shutil.rmtree(shard_dir)
    os.makedirs(shard_dir)

    index_rows = []
    for i, row in enumerate(corpus_rows):
        fname = slugify(i, row["title"]) + ".txt"
        with open(os.path.join(shard_dir, fname), "w", encoding="utf-8") as f:
            f.write(row["content"])
        index_rows.append({
            "title": row["title"],
            "domain": row["domain"],
            "ext": row["ext"],
            "size": row["size"],
            "n_questions": row["n_questions"],
            "file": f"corpus_{set_name}/{fname}",
        })

    corpus_path = os.path.join(ROOT, f"corpus_{set_name}.json")
    eval_path = os.path.join(ROOT, f"eval_{set_name}.json")
    with open(corpus_path, "w", encoding="utf-8") as f:
        json.dump(index_rows, f, ensure_ascii=False)
    with open(eval_path, "w", encoding="utf-8") as f:
        json.dump(eval_rows, f, ensure_ascii=False)

    domains = sorted({d["domain"] for d in eval_rows})
    corpus_mb = sum(r["size"] for r in corpus_rows) / 1e6
    upsert_manifest(set_name, {
        "set": set_name,
        "n_docs": len(index_rows),
        "n_questions": len(eval_rows),
        "domains": domains,
        "corpus_file": f"corpus_{set_name}.json",
        "eval_file": f"eval_{set_name}.json",
    })

    print(f"[{set_name}] records={n_records} questions={len(eval_rows)} "
          f"unique_docs={len(index_rows)} domains={domains}")
    print(f"[{set_name}] corpus text={corpus_mb:.2f}MB across {len(index_rows)} "
          f"shards · index={os.path.getsize(corpus_path)/1e6:.2f}MB · "
          f"eval={os.path.getsize(eval_path)/1e6:.2f}MB")
    if n_mismatch:
        print(f"[{set_name}] WARNING: {n_mismatch} records had "
              f"doc_files/# Document marker count mismatch")
    if n_varied:
        print(f"[{set_name}] note: {n_varied} doc occurrences varied after "
              f"question-block stripping (kept longest)")


def upsert_manifest(set_name, entry):
    path = os.path.join(ROOT, "sets.json")
    sets = []
    if os.path.exists(path):
        with open(path, encoding="utf-8") as f:
            sets = json.load(f)
    sets = [s for s in sets if s.get("set") != set_name]
    sets.append(entry)
    # Order sets by their numeric context length (128k < 1m < 4m < 10m).
    def sort_key(s):
        m = re.match(r"([0-9.]+)\s*([kmg]?)", s["set"].lower())
        if not m:
            return 0.0
        num = float(m.group(1) or 0)
        mult = {"": 1, "k": 1e3, "m": 1e6, "g": 1e9}.get(m.group(2), 1)
        return num * mult
    sets.sort(key=sort_key)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(sets, f, ensure_ascii=False, indent=2)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--set", required=True, help="set label, e.g. 128k or 1m")
    ap.add_argument("--input", help="path to the raw <set>_4domains.jsonl")
    args = ap.parse_args()
    input_path = args.input or os.path.join(
        DEFAULT_DATA_DIR, f"{args.set}_4domains.jsonl")
    if not os.path.exists(input_path):
        raise SystemExit(f"input not found: {input_path}")
    build_set(args.set, input_path)


if __name__ == "__main__":
    main()