|
| 1 | +# SPDX-License-Identifier: PMPL-1.0-or-later |
| 2 | +# Copyright (c) 2026 Jonathan D.A. Jewell <j.d.a.jewell@open.ac.uk> |
| 3 | +# |
| 4 | +# corpus_loader.jl — Single authoritative pass over per-prover JSONL. |
| 5 | +# |
| 6 | +# All metrics share one in-memory index keyed by prover → Vector{Record}. |
| 7 | +# This avoids N× re-reading the 286K-proof corpus across eight metrics. |
| 8 | + |
| 9 | +module CorpusLoader |
| 10 | + |
| 11 | +using JSON3 |
| 12 | + |
| 13 | +export load_corpus, CorpusRecord, CorpusIndex, PER_PROVER_FILES |
| 14 | + |
| 15 | +# Same list merge_corpus.jl consumes. Intentionally excludes aggregates. |
| 16 | +const PER_PROVER_FILES = [ |
| 17 | + "proof_states_mathlib4_max.jsonl", |
| 18 | + "proof_states_coqgym_max.jsonl", |
| 19 | + "proof_states_smtlib.jsonl", |
| 20 | + "proof_states_metamath.jsonl", |
| 21 | + "proof_states_hol_light.jsonl", |
| 22 | + "proof_states_hol4.jsonl", |
| 23 | + "proof_states_acl2.jsonl", |
| 24 | + "proof_states_pvs.jsonl", |
| 25 | + "proof_states_why3.jsonl", |
| 26 | + "proof_states_dafny.jsonl", |
| 27 | + "proof_states_fstar.jsonl", |
| 28 | + "proof_states_idris2.jsonl", |
| 29 | + "proof_states_mizar.jsonl", |
| 30 | + "proof_states_nuprl.jsonl", |
| 31 | + "proof_states_minlog.jsonl", |
| 32 | + "proof_states_twelf.jsonl", |
| 33 | + "proof_states_imandra.jsonl", |
| 34 | + "proof_states_minizinc.jsonl", |
| 35 | + "proof_states_afp.jsonl", |
| 36 | + "proof_states_agda.jsonl", |
| 37 | + "proof_states_typechecker_ecosystem.jsonl", |
| 38 | +] |
| 39 | + |
| 40 | +struct CorpusRecord |
| 41 | + id::Int |
| 42 | + prover::String |
| 43 | + theorem::String |
| 44 | + goal::String |
| 45 | + context::Vector{String} |
| 46 | + tactic_proof::String |
| 47 | + source::String |
| 48 | +end |
| 49 | + |
| 50 | +const CorpusIndex = Dict{String, Vector{CorpusRecord}} |
| 51 | + |
| 52 | +function _string_vec(v)::Vector{String} |
| 53 | + v isa AbstractVector || return String[] |
| 54 | + out = String[] |
| 55 | + for x in v |
| 56 | + x isa AbstractString && push!(out, String(x)) |
| 57 | + end |
| 58 | + return out |
| 59 | +end |
| 60 | + |
| 61 | +function _coerce(rec::AbstractDict)::CorpusRecord |
| 62 | + CorpusRecord( |
| 63 | + Int(get(rec, "id", 0)), |
| 64 | + String(get(rec, "prover", "unknown")), |
| 65 | + String(get(rec, "theorem", "")), |
| 66 | + String(get(rec, "goal", "")), |
| 67 | + _string_vec(get(rec, "context", String[])), |
| 68 | + String(get(rec, "tactic_proof", "")), |
| 69 | + String(get(rec, "source", "")), |
| 70 | + ) |
| 71 | +end |
| 72 | + |
| 73 | +""" |
| 74 | + load_corpus(training_dir) -> CorpusIndex |
| 75 | +
|
| 76 | +Read every authoritative per-prover JSONL under `training_dir`. |
| 77 | +Returns a dict mapping prover name → records. Missing files are skipped. |
| 78 | +""" |
| 79 | +function load_corpus(training_dir::AbstractString)::CorpusIndex |
| 80 | + idx = CorpusIndex() |
| 81 | + for fname in PER_PROVER_FILES |
| 82 | + path = joinpath(training_dir, fname) |
| 83 | + isfile(path) || continue |
| 84 | + open(path, "r") do fh |
| 85 | + for line in eachline(fh) |
| 86 | + s = strip(line) |
| 87 | + isempty(s) && continue |
| 88 | + try |
| 89 | + rec = JSON3.read(s, Dict{String,Any}) |
| 90 | + cr = _coerce(rec) |
| 91 | + bucket = get!(idx, cr.prover, CorpusRecord[]) |
| 92 | + push!(bucket, cr) |
| 93 | + catch |
| 94 | + end |
| 95 | + end |
| 96 | + end |
| 97 | + end |
| 98 | + return idx |
| 99 | +end |
| 100 | + |
| 101 | +end |
0 commit comments