lean-bench/textbook/extract_exercises.py

"""Slice the 'Exercises N.N' block out of a pdftotext -layout dump and list its
numbered problems, as input for formalize_exercises.py.

Mirrors extract_section.py but for the exercise list at the end of a section, which
uses bare 'N.N.N.' numbering at the start of a line (no Theorem/Proposition keyword)
instead of labeled Theorem/Proposition/Lemma/Corollary/Definition entries. Exercise
numbers are their own independent sequence (Exercise 1.6.5 has nothing to do with
Proposition/Theorem/Lemma/Corollary/Definition 1.6.5 in the body) -- keep their ids
namespaced separately downstream so the two never collide.

A known false-positive source: a sentence like "...converse of Proposition\n1.6.19."
can wrap so the bare number lands alone at the start of a line, which looks
identical to a real exercise heading. numbered_items() drops any item whose number
jumps ahead of the running sequence by more than --max-gap AND has no text on its
own line (a real item almost always has problem text right after the number; a
wrapped cross-reference is just the number and a period). Dropped candidates are
still reported, never silently discarded.
"""
import argparse
import json
import re
import sys
from pathlib import Path

ITEM_RE = re.compile(r'^(\d+\.\d+\.\d+)\.\s*(.*)$')


def find_heading_line(lines, heading):
    for i, line in enumerate(lines):
        if line.strip() == heading:
            return i
    raise ValueError(f'Heading not found: {heading!r}')


def extract_exercises(full_text, start_heading, end_heading):
    lines = full_text.split('\n')
    start = find_heading_line(lines, start_heading)
    end = find_heading_line(lines, end_heading) if end_heading else len(lines)
    if end <= start:
        raise ValueError('end_heading must come after start_heading')
    return '\n'.join(lines[start:end])


def numbered_items(exercises_text, section_prefix, max_gap=3):
    """Unique numbered items in exercises_text whose number starts with
    section_prefix (e.g. '1.6.'), keeping only the first occurrence of each number
    (a running page header can repeat text near a page break). Returned in book
    order (numeric, not lexical). See module docstring for the wrapped-cross-
    reference false-positive filter; dropped candidates are returned separately so
    callers can report them instead of silently losing items."""
    seen = {}
    for line in exercises_text.split('\n'):
        line = line.strip()
        m = ITEM_RE.match(line)
        if not m:
            continue
        number, rest = m.group(1), m.group(2)
        if not number.startswith(section_prefix) or number in seen:
            continue
        seen[number] = {"number": number, "first_line": line[:120], "rest_empty": not rest.strip()}

    def sort_key(item):
        return tuple(int(p) for p in item['number'].split('.'))

    ordered = sorted(seen.values(), key=sort_key)
    kept, dropped = [], []
    last_n = 0
    for item in ordered:
        n = int(item['number'].split('.')[-1])
        if item['rest_empty'] and last_n and (n - last_n) > max_gap:
            dropped.append(item)
            continue
        kept.append({"number": item['number'], "first_line": item['first_line']})
        last_n = n
    return kept, dropped


def main():
    parser = argparse.ArgumentParser(description='Slice one "Exercises N.N" block out of a pdftotext -layout dump')
    parser.add_argument('layout_txt', type=Path)
    parser.add_argument('--start-heading', required=True, help="e.g. 'Exercises 1.6'")
    parser.add_argument('--end-heading', help="e.g. '1.7. Modular Arithmetic'; omit for end of file")
    parser.add_argument('--section-prefix', required=True, help="e.g. '1.6.'")
    parser.add_argument('--max-gap', type=int, default=3, help='see module docstring')
    parser.add_argument('--out-text', type=Path, required=True)
    parser.add_argument('--out-index', type=Path, required=True)
    args = parser.parse_args()

    full_text = args.layout_txt.read_text(encoding='utf-8')
    exercises_text = extract_exercises(full_text, args.start_heading, args.end_heading)
    args.out_text.write_text(exercises_text, encoding='utf-8')
    index, dropped = numbered_items(exercises_text, args.section_prefix, args.max_gap)
    args.out_index.write_text(json.dumps(index, ensure_ascii=False, indent=2) + '\n', encoding='utf-8')
    print(f'{len(index)} numbered exercises -> {args.out_index}')
    for item in dropped:
        print(f"  dropped as likely a wrapped cross-reference, not a real exercise: {item['number']}"
              f" ({item['first_line']!r})", file=sys.stderr)


if __name__ == '__main__':
    main()