|
| 1 | +#!/usr/bin/env julia |
| 2 | +# SPDX-License-Identifier: PMPL-1.0-or-later |
| 3 | +# SPDX-FileCopyrightText: 2026 Jonathan D.A. Jewell (hyperpolymath) <j.d.a.jewell@open.ac.uk> |
| 4 | +# |
| 5 | +# align_premises.jl — rebuild training_data/premises_COMPLETE.jsonl with proof_ids |
| 6 | +# that match the fresh sequential ids written by merge_corpus.jl into |
| 7 | +# training_data/proof_states_UNIFIED.jsonl. |
| 8 | +# |
| 9 | +# Background |
| 10 | +# ---------- |
| 11 | +# The per-prover extractors (scripts/extract_*.jl and scripts/max_extract_*.jl) |
| 12 | +# emit proof_states and premises sharing a per-extractor-local integer id space: |
| 13 | +# |
| 14 | +# {"id": 60123, "prover": "Lean", "theorem": "Foo.bar", ...} (proof_state) |
| 15 | +# {"proof_id": 60123, "prover": "Lean", "theorem": "Foo.bar", ...} (premise) |
| 16 | +# |
| 17 | +# scripts/merge_corpus.jl then writes proof_states_UNIFIED.jsonl after dedup |
| 18 | +# by (prover, theorem) and **rewrites every proof_state's id to a fresh |
| 19 | +# sequential counter (1..N)**. The premise files retain the original |
| 20 | +# extractor ids, so proof_state.id and premise.proof_id no longer share an id |
| 21 | +# space — the dataloader's join by proof_id matches ~0% of records. |
| 22 | +# |
| 23 | +# This script rebuilds premises_COMPLETE.jsonl using the |
| 24 | +# (canonical_prover, theorem) pair as the durable join key (which merge_corpus |
| 25 | +# already guarantees is unique within UNIFIED). For every premise record, |
| 26 | +# we look up the fresh id assigned by merge_corpus and rewrite |
| 27 | +# `proof_id := UNIFIED_id`. Records whose (prover, theorem) does not appear |
| 28 | +# in UNIFIED are dropped and counted. |
| 29 | +# |
| 30 | +# Usage |
| 31 | +# ----- |
| 32 | +# julia --project=src/julia scripts/align_premises.jl |
| 33 | +# |
| 34 | +# Outputs |
| 35 | +# ------- |
| 36 | +# training_data/premises_COMPLETE.jsonl — aligned premises |
| 37 | +# training_data/premises_ALIGNED_stats.json — coverage report |
| 38 | +# |
| 39 | +# Run this after merge_corpus.jl and before run_training.jl. |
| 40 | + |
| 41 | +using JSON3 |
| 42 | + |
| 43 | +# --------------------------------------------------------------------------- |
| 44 | +# Paths |
| 45 | +# --------------------------------------------------------------------------- |
| 46 | + |
| 47 | +const REPO_ROOT = dirname(dirname(abspath(@__FILE__))) |
| 48 | +const TRAINING_DIR = joinpath(REPO_ROOT, "training_data") |
| 49 | + |
| 50 | +const UNIFIED_PATH = joinpath(TRAINING_DIR, "proof_states_UNIFIED.jsonl") |
| 51 | +const OUT_PATH = joinpath(TRAINING_DIR, "premises_COMPLETE.jsonl") |
| 52 | +const STATS_PATH = joinpath(TRAINING_DIR, "premises_ALIGNED_stats.json") |
| 53 | + |
| 54 | +# Every premises_*.jsonl in training_data/ is a candidate source, except the |
| 55 | +# output file itself and any ALIGNED_stats JSON we emit. We also skip the |
| 56 | +# timestamped `premises_<YYYY-MM-DD>.jsonl` snapshots by default but include |
| 57 | +# them when ECHIDNA_ALIGN_INCLUDE_DATED=1. |
| 58 | +function collect_premise_sources() |
| 59 | + sources = String[] |
| 60 | + for f in readdir(TRAINING_DIR; join=true) |
| 61 | + base = basename(f) |
| 62 | + startswith(base, "premises_") || continue |
| 63 | + endswith(base, ".jsonl") || continue |
| 64 | + base == "premises_COMPLETE.jsonl" && continue |
| 65 | + if get(ENV, "ECHIDNA_ALIGN_INCLUDE_DATED", "0") != "1" && |
| 66 | + occursin(r"^premises_\d{4}-\d{2}-\d{2}\.jsonl$", base) |
| 67 | + continue |
| 68 | + end |
| 69 | + push!(sources, f) |
| 70 | + end |
| 71 | + sort!(sources) |
| 72 | + return sources |
| 73 | +end |
| 74 | + |
| 75 | +# --------------------------------------------------------------------------- |
| 76 | +# Prover-name canonicalisation — must match merge_corpus.jl's mapping. |
| 77 | +# Keep this in lockstep with scripts/merge_corpus.jl's PROVER_CANONICAL. |
| 78 | +# --------------------------------------------------------------------------- |
| 79 | + |
| 80 | +const PROVER_CANONICAL = Dict( |
| 81 | + "Lean4" => "Lean", |
| 82 | + "Rocq" => "Coq", |
| 83 | + "Fstar" => "F*", |
| 84 | + "CVC4" => "CVC5", |
| 85 | + "AltErgo" => "Alt-Ergo", |
| 86 | + "HOL Light" => "HOLLight", |
| 87 | + "E Prover" => "EProver", |
| 88 | + "OR-Tools" => "ORTools", |
| 89 | +) |
| 90 | + |
| 91 | +canon(prover::AbstractString) = get(PROVER_CANONICAL, String(prover), String(prover)) |
| 92 | + |
| 93 | +# --------------------------------------------------------------------------- |
| 94 | +# Build the (canonical_prover, theorem) → fresh_id index from UNIFIED. |
| 95 | +# --------------------------------------------------------------------------- |
| 96 | + |
| 97 | +function build_unified_index(path::String) |
| 98 | + isfile(path) || error("missing $(path); run scripts/merge_corpus.jl first.") |
| 99 | + index = Dict{Tuple{String,String},Int}() |
| 100 | + open(path, "r") do fh |
| 101 | + for line in eachline(fh) |
| 102 | + isempty(strip(line)) && continue |
| 103 | + entry = try |
| 104 | + JSON3.read(line, Dict{String,Any}) |
| 105 | + catch |
| 106 | + continue |
| 107 | + end |
| 108 | + prover = canon(get(entry, "prover", "")) |
| 109 | + theorem = String(get(entry, "theorem", "")) |
| 110 | + new_id = get(entry, "id", nothing) |
| 111 | + (isempty(prover) || isempty(theorem) || new_id === nothing) && continue |
| 112 | + key = (prover, theorem) |
| 113 | + # First-write wins. merge_corpus already dedupes by (prover,theorem) |
| 114 | + # so collisions here would indicate a merge bug — report and keep |
| 115 | + # the earlier id. |
| 116 | + if haskey(index, key) && index[key] != new_id |
| 117 | + @warn "duplicate (prover,theorem) in UNIFIED" prover=prover theorem=theorem first=index[key] second=new_id |
| 118 | + continue |
| 119 | + end |
| 120 | + index[key] = Int(new_id) |
| 121 | + end |
| 122 | + end |
| 123 | + return index |
| 124 | +end |
| 125 | + |
| 126 | +# --------------------------------------------------------------------------- |
| 127 | +# Walk every premise source, rewrite proof_id, and emit premises_COMPLETE. |
| 128 | +# --------------------------------------------------------------------------- |
| 129 | + |
| 130 | +function align_premises(index::Dict{Tuple{String,String},Int}, sources::Vector{String}) |
| 131 | + written = 0 |
| 132 | + dropped_no_theorem = 0 |
| 133 | + dropped_no_match = 0 |
| 134 | + dropped_malformed = 0 |
| 135 | + per_source = Dict{String,Dict{String,Int}}() |
| 136 | + |
| 137 | + open(OUT_PATH, "w") do out |
| 138 | + for src in sources |
| 139 | + s_written = 0 |
| 140 | + s_dropped = 0 |
| 141 | + open(src, "r") do fh |
| 142 | + for line in eachline(fh) |
| 143 | + isempty(strip(line)) && continue |
| 144 | + rec = try |
| 145 | + JSON3.read(line, Dict{String,Any}) |
| 146 | + catch |
| 147 | + dropped_malformed += 1 |
| 148 | + s_dropped += 1 |
| 149 | + continue |
| 150 | + end |
| 151 | + prover = canon(get(rec, "prover", "")) |
| 152 | + theorem = String(get(rec, "theorem", "")) |
| 153 | + if isempty(prover) || isempty(theorem) |
| 154 | + dropped_no_theorem += 1 |
| 155 | + s_dropped += 1 |
| 156 | + continue |
| 157 | + end |
| 158 | + new_id = get(index, (prover, theorem), nothing) |
| 159 | + if new_id === nothing |
| 160 | + dropped_no_match += 1 |
| 161 | + s_dropped += 1 |
| 162 | + continue |
| 163 | + end |
| 164 | + rec["proof_id"] = new_id |
| 165 | + println(out, JSON3.write(rec)) |
| 166 | + written += 1 |
| 167 | + s_written += 1 |
| 168 | + end |
| 169 | + end |
| 170 | + per_source[basename(src)] = Dict( |
| 171 | + "written" => s_written, |
| 172 | + "dropped" => s_dropped, |
| 173 | + ) |
| 174 | + end |
| 175 | + end |
| 176 | + |
| 177 | + return ( |
| 178 | + written = written, |
| 179 | + dropped_no_theorem = dropped_no_theorem, |
| 180 | + dropped_no_match = dropped_no_match, |
| 181 | + dropped_malformed = dropped_malformed, |
| 182 | + per_source = per_source, |
| 183 | + ) |
| 184 | +end |
| 185 | + |
| 186 | +# --------------------------------------------------------------------------- |
| 187 | +# Main |
| 188 | +# --------------------------------------------------------------------------- |
| 189 | + |
| 190 | +function main() |
| 191 | + println("align_premises — reconcile premise proof_ids with proof_states_UNIFIED") |
| 192 | + println("=" ^ 72) |
| 193 | + |
| 194 | + println("\nBuilding (prover, theorem) → fresh_id index from UNIFIED...") |
| 195 | + index = build_unified_index(UNIFIED_PATH) |
| 196 | + println(" $(length(index)) (prover, theorem) pairs loaded.") |
| 197 | + |
| 198 | + sources = collect_premise_sources() |
| 199 | + println("\nPremise sources ($(length(sources)) files):") |
| 200 | + for s in sources |
| 201 | + println(" $(basename(s))") |
| 202 | + end |
| 203 | + |
| 204 | + println("\nAligning...") |
| 205 | + result = align_premises(index, sources) |
| 206 | + |
| 207 | + total_in = result.written + result.dropped_no_theorem + |
| 208 | + result.dropped_no_match + result.dropped_malformed |
| 209 | + match_rate = total_in == 0 ? 0.0 : 100.0 * result.written / total_in |
| 210 | + |
| 211 | + println("\nResults") |
| 212 | + println("-" ^ 72) |
| 213 | + println(" Wrote : $(result.written) records to $(OUT_PATH)") |
| 214 | + println(" Dropped (no key) : $(result.dropped_no_theorem) (missing prover or theorem)") |
| 215 | + println(" Dropped (no match): $(result.dropped_no_match) (theorem not in UNIFIED)") |
| 216 | + println(" Dropped (bad json): $(result.dropped_malformed)") |
| 217 | + println(" Match rate : $(round(match_rate; digits=2))%") |
| 218 | + |
| 219 | + println("\nPer source:") |
| 220 | + for (name, counts) in sort(collect(result.per_source); by = p -> p[1]) |
| 221 | + println(" $(rpad(name, 40)) written=$(counts["written"]) dropped=$(counts["dropped"])") |
| 222 | + end |
| 223 | + |
| 224 | + # Persist a machine-readable stats file alongside the output. |
| 225 | + stats = Dict( |
| 226 | + "unified_path" => UNIFIED_PATH, |
| 227 | + "output_path" => OUT_PATH, |
| 228 | + "unified_index_size" => length(index), |
| 229 | + "sources" => [basename(s) for s in sources], |
| 230 | + "written" => result.written, |
| 231 | + "dropped_no_theorem" => result.dropped_no_theorem, |
| 232 | + "dropped_no_match" => result.dropped_no_match, |
| 233 | + "dropped_malformed" => result.dropped_malformed, |
| 234 | + "match_rate_percent" => match_rate, |
| 235 | + "per_source" => result.per_source, |
| 236 | + ) |
| 237 | + open(STATS_PATH, "w") do fh |
| 238 | + JSON3.pretty(fh, stats) |
| 239 | + end |
| 240 | + println("\nStats written to $(STATS_PATH)") |
| 241 | + |
| 242 | + if result.written == 0 |
| 243 | + println("\nERROR: 0 records written. Check that the premise files carry 'prover' and 'theorem' fields.") |
| 244 | + exit(2) |
| 245 | + end |
| 246 | + return 0 |
| 247 | +end |
| 248 | + |
| 249 | +main() |
0 commit comments