spb/uqo-research-db

1# Batched full-corpus LLM extraction driver (v2, 2026-10-27).2# Sections are loaded per CIK-batch (bounds memory) and every processed section is3# recorded in the `extracted_sections` ledger, so a restart never re-pays for done4# work and newly parsed sections of already-covered companies are picked up.5#6#   PYTHONPATH=. ../.venv/bin/python run_full_extract.py [--workers 24] [--batch 20] [--dry-run]7from __future__ import annotations89import argparse10import sys1112from spid.config import get_logger13from spid.pipeline import run_extraction14from spid.store import Store1516log = get_logger("driver")171819def main() -> int:20    ap = argparse.ArgumentParser()21    ap.add_argument("--batch", type=int, default=20, help="CIKs per batch")22    ap.add_argument("--workers", type=int, default=24)23    ap.add_argument("--backend", default="llm")24    ap.add_argument("--max-sections", type=int, default=None, help="stop after roughly this many sections (budget cap)")25    ap.add_argument("--dry-run", action="store_true", help="only count pending sections per CIK")26    args = ap.parse_args()2728    store = Store()29    store.con.execute("SET arrow_large_buffer_size=true")30    pending = store.con.execute("""31        SELECT cik, COUNT(*) FROM v_sections32        WHERE word_count >= 30 AND section_id NOT IN (SELECT section_id FROM extracted_sections)33        GROUP BY cik ORDER BY cik""").fetchall()34    total = sum(n for _, n in pending)35    log.info("pending: %d sections over %d CIKs | batch=%d workers=%d backend=%s",36             total, len(pending), args.batch, args.workers, args.backend)37    if args.dry_run:38        return 03940    ciks = [c for c, _ in pending]41    n_batches = (len(ciks) + args.batch - 1) // args.batch42    cum_sec = cum_men = 043    for bi in range(n_batches):44        batch = ciks[bi * args.batch:(bi + 1) * args.batch]45        try:46            res = run_extraction(store, backend=args.backend, ciks=batch, forms=None, workers=args.workers)47        except Exception as e:  # noqa: BLE00148            log.error("BATCH %d FAILED (%s) ciks=%s - continuing", bi + 1, e, batch)49            continue50        cum_sec += res["sections"]51        cum_men += res["mentions"]52        log.info("BATCH %d/%d done | batch_sec=%d batch_men=%d | CUM sections=%d/%d mentions=%d",53                 bi + 1, n_batches, res["sections"], res["mentions"], cum_sec, total, cum_men)54        if args.max_sections and cum_sec >= args.max_sections:55            log.info("budget cap reached (%d sections) - stopping", cum_sec)56            break57    log.info("EXTRACT RUN COMPLETE | sections=%d mentions=%d | TOTAL mentions in DB=%d",58             cum_sec, cum_men, store.count("project_mentions"))59    store.close()60    return 0616263if __name__ == "__main__":64    sys.exit(main())65