scripts/build_examples.py

"""Build examples/examples.json for the demo.

Sources
- HotpotQA / MuSiQue / NQ samples produced by the upstream dataset providers
  (`lensvlm.dataset_providers`, 3K-12K token contexts with distractors and
  character-level evidence spans), exported to /tmp/exprep/*.json.
- The official `examples/hotpotqa_demo.json` from apple-aiml-research/ml-lensvlm.
- Two public-domain texts from Project Gutenberg: the U.S. Constitution (#5)
  and "A Scandal in Bohemia" from The Adventures of Sherlock Holmes (#1661).

Evidence spans for the non-benchmark documents are located by exact phrase
match; the app renders them with `render_pages(..., evidence_spans=...)` to
mark the gold pages at the selected compression.

Usage: python scripts/build_examples.py [candidates|final]

`candidates` writes every candidate; `final` keeps only the ids listed in
examples/selection.json (the candidates that were answered correctly, from
the expanded pages, on the hosted Space).
"""

import json
import os
import re
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PREP = "/tmp/exprep"
OUT = os.path.join(ROOT, "examples", "examples.json")


def sentence_span(text, phrase):
    """(start, end) of the sentence containing `phrase`."""
    i = text.index(phrase)
    start = max(text.rfind(". ", 0, i) + 2, text.rfind("\n", 0, i) + 1, 0)
    end = text.find(". ", i + len(phrase))
    end = len(text) if end == -1 else end + 1
    return [start, end]


def phrase_span(text, phrase):
    i = text.index(phrase)
    return [i, i + len(phrase)]


def benchmark(name, idx, compression, label, phrases=None):
    """`phrases` overrides the provider's evidence spans with the sentences
    containing them (NQ marks the whole Wikipedia page as evidence)."""
    s = json.load(open(os.path.join(PREP, f"{name}_validation.json")))[idx]
    spans = ([sentence_span(s["context"], p) for p in phrases] if phrases
             else [list(sp) for sp in s["evidence_spans"]])
    return dict(
        id=s["id"],
        label=label,
        source={"hotpotqa": "HotpotQA (distractor)", "musique": "MuSiQue",
                "nq": "Natural Questions"}[name],
        question=s["question"],
        context=s["context"],
        compression=compression,
        answers=list(dict.fromkeys(s["answers"] or [s["answer"]])),
        match_all=False,
        evidence_spans=spans,
    )


def official_demo():
    d = json.load(open("/tmp/ml-lensvlm/examples/hotpotqa_demo.json"))
    ctx = d["context"]
    return dict(
        id="hotpotqa_demo",
        label="Official demo: Kiss and Tell → government post (HotpotQA, 10×)",
        source="HotpotQA (official LensVLM demo sample)",
        question=d["question"],
        context=ctx,
        compression="10x",
        answers=d["answers"],
        match_all=False,
        evidence_spans=[
            sentence_span(ctx, "Kiss and Tell is a 1945 American comedy film starring"),
            sentence_span(ctx, "also served as Chief of Protocol of the United States"),
        ],
    )


def constitution(question, answers, phrases, compression, label, match_all=False):
    raw = open(os.path.join(PREP, "constitution.txt"), encoding="utf-8").read()
    ctx = raw[raw.index("THE CONSTITUTION OF THE UNITED STATES OF AMERICA, 1787"):raw.index("*** END")].strip()
    return dict(
        id="constitution_" + re.sub(r"\W+", "_", answers[0].split("|")[0].lower()),
        label=label,
        source="U.S. Constitution (Project Gutenberg eBook #5, public domain)",
        question=question,
        context=ctx,
        compression=compression,
        answers=answers,
        match_all=match_all,
        evidence_spans=[phrase_span(ctx, p) for p in phrases],
    )


def holmes(question, answers, phrases, compression, label, match_all=False):
    lines = open(os.path.join(PREP, "holmes.txt"), encoding="utf-8").read().split("\n")
    ctx = "\n".join(lines[54:1135])
    ctx = ctx[:ctx.index("II. THE RED-HEADED LEAGUE")].strip()
    assert ctx.startswith("I. A SCANDAL IN BOHEMIA") and ctx.endswith("_the_ woman."), ctx[-50:]
    return dict(
        id="holmes_" + re.sub(r"\W+", "_", answers[0].split("|")[0].lower()),
        label=label,
        source="A Scandal in Bohemia, A. Conan Doyle (Project Gutenberg eBook #1661, public domain)",
        question=question,
        context=ctx,
        compression=compression,
        answers=answers,
        match_all=match_all,
        evidence_spans=[phrase_span(ctx, p) for p in phrases],
    )


def candidates():
    return [
        official_demo(),
        benchmark("hotpotqa", 5, "10x", "HotpotQA: Danny DeVito film scored by Kool Kojak (10×)"),
        benchmark("hotpotqa", 8, "10x", "HotpotQA: band named by drawing from a hat (10×)"),
        benchmark("hotpotqa", 13, "15x", "HotpotQA: port city north of Lingnan Fine Arts Museum (15×)"),
        benchmark("hotpotqa", 2, "15x", "HotpotQA: Art Laboe & Scout Tufankjian's ancestry (15×)"),
        benchmark("musique", 0, "10x", "MuSiQue: Thelma's actress in A League of Their Own (10×)"),
        benchmark("musique", 7, "15x", "MuSiQue: spouse of The Mighty B's creator (15×)"),
        benchmark("musique", 26, "10x", "MuSiQue: mother of Julia Murdock Smith's father (10×)"),
        benchmark("nq", 21, "15x", "Natural Questions: who wrote City of New Orleans (15×)",
                  ["is a folk song written by Steve Goodman"]),
        benchmark("nq", 17, "15x", "Natural Questions: when UK national service ended (15×)",
                  ["call - ups formally ended on 31 December 1960"]),
        constitution(
            "Which body has the power to declare war, and who is Commander in Chief of the Army and Navy?",
            ["Congress", "President"],
            ["To declare War", "The President shall be Commander in Chief"],
            "10x", "U.S. Constitution: war powers vs. Commander in Chief (10×)", match_all=True,
        ),
        constitution(
            "Who attested the signatures on the Constitution, and in what role?",
            ["William Jackson", "Secretary"],
            ["William Jackson, Secretary"],
            "10x", "U.S. Constitution: who attested the signatures (10×)", match_all=True,
        ),
        holmes(
            "Under what name did the King of Bohemia first introduce himself to Holmes?",
            ["Von Kramm"],
            ["You may address me as the Count Von Kramm"],
            "10x", "Sherlock Holmes: the King's alias (full story, 10×)",
        ),
        holmes(
            "Where exactly had Irene Adler hidden the photograph in her sitting-room?",
            ["bell-pull"],
            ["photograph is in a recess behind a sliding panel just above the right\nbell-pull"],
            "10x", "Sherlock Holmes: where the photograph was hidden (full story, 10×)",
        ),
        holmes(
            "In which church were Irene Adler and Godfrey Norton married?",
            ["St. Monica"],
            ["the Church of St. Monica in the\nEdgeware Road"],
            "15x", "Sherlock Holmes: the wedding church (full story, 15×)",
        ),
        # Second screening round.
        benchmark("hotpotqa", 9, "10x", "HotpotQA: birth year of the singer of “Tunak” (10×)"),
        benchmark("hotpotqa", 27, "10x", "HotpotQA: birth year of Lisa Simpson's creator (10×)"),
        benchmark("hotpotqa", 37, "10x", "HotpotQA: film by Joby Harold scored by Samuel Sim (10×)"),
        benchmark("hotpotqa", 26, "15x", "HotpotQA: home city of Ruddles Brewery's owner (15×)"),
        benchmark("musique", 12, "10x", "MuSiQue: father of Crockett County's namesake (10×)"),
        benchmark("musique", 15, "10x", "MuSiQue: birthplace of the president born on July 4 (10×)"),
        benchmark("musique", 36, "15x", "MuSiQue: where the first female Justice studied (15×)"),
        constitution(
            "What minimum ages does the Constitution set for a Representative, a Senator and the President?",
            ["twenty five|25", "thirty Years|30", "thirty five|35"],
            ["Age of twenty five Years", "Age of\nthirty Years", "Age of thirty five Years"],
            "10x", "U.S. Constitution: minimum ages for Congress and President (10×)", match_all=True,
        ),
        constitution(
            "How old must a person be to be eligible to the Office of President, and how long must they have been a resident of the United States?",
            ["thirty five|35", "fourteen|14"],
            ["Age of thirty five Years,\nand been fourteen Years a Resident"],
            "10x", "U.S. Constitution: eligibility for President (10×)", match_all=True,
        ),
        holmes(
            "How much money did the King leave with Holmes for present expenses?",
            ["three hundred pounds|300", "seven hundred|700"],
            ["three hundred pounds in gold and seven hundred in notes"],
            # at 10x the first page pick changed between deployments of identical code (page 9, later 19);
            # 5x read the evidence page on both boots checked since
            "5x", "Sherlock Holmes: the King's expense money (full story, 5×)", match_all=True,
        ),
        holmes(
            "In what disguise did Holmes get himself into Briony Lodge?",
            ["clergyman"],
            ["Nonconformist clergyman"],
            "10x", "Sherlock Holmes: Holmes's disguise (full story, 10×)",
        ),
        holmes(
            "Who was Irene Adler's frequent male visitor, and where was he from?",
            ["Godfrey Norton", "Inner Temple"],
            ["Mr. Godfrey\nNorton, of the Inner Temple"],
            "10x", "Sherlock Holmes: Irene Adler's visitor (full story, 10×)", match_all=True,
        ),
    ]


if __name__ == "__main__":
    mode = sys.argv[1] if len(sys.argv) > 1 else "candidates"
    exs = candidates()
    if mode == "final":
        keep = json.load(open(os.path.join(ROOT, "examples", "selection.json")))["keep"]
        by_id = {e["id"]: e for e in exs}
        exs = [by_id[k] for k in keep]
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    json.dump(exs, open(OUT, "w"), ensure_ascii=False, indent=1)
    for e in exs:
        print(f"{e['id']:32s} {e['compression']:4s} chars={len(e['context']):6d} spans={len(e['evidence_spans'])}  {e['question'][:70]}")
    print(len(exs), "examples ->", OUT)