Skip to content

Commit 1d9a76a

Browse files
Claude/amazing cori kw a2 v (#101)
Co-authored-by: Claude <noreply@anthropic.com>
1 parent 946541f commit 1d9a76a

9 files changed

Lines changed: 187 additions & 14 deletions

Justfile

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,18 @@ retrain: align-premises
337337
retrain-skip-align:
338338
julia --project=src/julia src/julia/run_training.jl
339339

340+
# Full training run on the real corpus (auto-selects GPU when available).
341+
# Honours ECHIDNA_MAX_PROOF_STATES, ECHIDNA_NUM_EPOCHS, ECHIDNA_NUM_NEGATIVES.
342+
# Produces models/neural/{premise_selector,tactic_predictor}.bson, vocab.json,
343+
# updates models/model_metadata.txt, and appends to training_data/metrics_baseline.jsonl.
344+
train:
345+
julia --project=src/julia src/julia/run_training.jl
346+
347+
# CPU smoke pass: 2000 proof states, 2 epochs — fast end-to-end sanity check.
348+
# Safe to run on any dev box without a GPU. Produces the same artefacts as `train`.
349+
train-cpu:
350+
ECHIDNA_MAX_PROOF_STATES=2000 ECHIDNA_NUM_EPOCHS=2 julia --project=src/julia src/julia/run_training_cpu.jl
351+
340352
# End-to-end pipeline: provision → extract → merge → align → retrain.
341353
# Use `ECHIDNA_MAX_PROOF_STATES=0 just corpus-refresh` to lift the sample cap.
342354
corpus-refresh: provision-corpora extract-corpora merge-corpora align-premises retrain

src/julia/api/gnn_endpoint.jl

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -42,27 +42,37 @@ const TOTAL_TRAINING_RECORDS = Ref{Int}(0)
4242
"""
4343
load_gnn_model(models_dir::String)
4444
45-
Load the GNN premise ranker model from disk.
46-
Falls back to creating a fresh (untrained) model if no checkpoint exists.
45+
Load the GNN premise ranker model from disk. Tries, in order:
46+
1. `models/neural/gnn_ranker/` — directory written by run_training*.jl
47+
2. `models/neural/best_model/` — early-stopping checkpoint
48+
3. `models/neural/final_model/` — last epoch checkpoint
49+
50+
Only falls back to the cosine path (GNN_MODEL[] = nothing) when none of
51+
the above directories exist or all fail to deserialise. The cosine path
52+
is the genuine missing-model fallback for CI smoke runs.
4753
"""
4854
function load_gnn_model(models_dir::String)
49-
model_path = joinpath(models_dir, "neural", "gnn_ranker")
50-
51-
if isdir(model_path)
52-
@info "Loading GNN model from $model_path"
55+
candidate_dirs = [
56+
joinpath(models_dir, "neural", "gnn_ranker"),
57+
joinpath(models_dir, "neural", "best_model"),
58+
joinpath(models_dir, "neural", "final_model"),
59+
]
60+
61+
for model_path in candidate_dirs
62+
isdir(model_path) || continue
63+
@info "Trying to load GNN model from $model_path"
5364
try
5465
solver = load_solver(model_path)
5566
GNN_MODEL[] = solver
56-
@info "GNN model loaded successfully"
67+
@info "GNN model loaded successfully from $model_path"
5768
return true
5869
catch e
59-
@warn "Failed to load GNN model: $e"
70+
@warn "Failed to load GNN model from $model_path: $e"
6071
end
6172
end
6273

63-
@info "No trained GNN model found — creating fresh model for inference"
64-
# Create a minimal model with default configuration for the endpoint
65-
# to respond (scores will be random until training completes)
74+
# All candidates exhausted — cosine fallback is genuine missing-model path.
75+
@warn "No trained GNN model found in $(joinpath(models_dir, "neural")) — ranking will use cosine similarity until weights are trained (run: just train-cpu)"
6676
GNN_MODEL[] = nothing
6777
return false
6878
end

src/julia/train_advanced_models.jl renamed to src/julia/archive/train_advanced_models.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env julia
2+
# ARCHIVED 2026-05-24 — superseded by src/julia/run_training.jl. Kept for history only; do not invoke.
23
# SPDX-FileCopyrightText: 2026 ECHIDNA Project Team
34
# SPDX-License-Identifier: MPL-2.0
45

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env julia
2+
# ARCHIVED 2026-05-24 — superseded by src/julia/run_training.jl. Kept for history only; do not invoke.
23
# SPDX-License-Identifier: MPL-2.0
34
# Train GNN model and evaluate on validation set, outputting metrics for health monitoring
45

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env julia
2+
# ARCHIVED 2026-05-24 — superseded by src/julia/run_training.jl. Kept for history only; do not invoke.
23
# SPDX-FileCopyrightText: 2026 ECHIDNA Project Team
34
# SPDX-License-Identifier: MPL-2.0
45

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env julia
2+
# ARCHIVED 2026-05-24 — superseded by src/julia/run_training.jl. Kept for history only; do not invoke.
23
# SPDX-FileCopyrightText: 2026 ECHIDNA Project Team
34
# SPDX-License-Identifier: MPL-2.0
45

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
# ARCHIVED 2026-05-24 — superseded by src/julia/run_training.jl. Kept for history only; do not invoke.
12
# SPDX-FileCopyrightText: 2026 ECHIDNA Project Team
23
# SPDX-License-Identifier: MPL-2.0
34

src/julia/run_training.jl

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727

2828
using Pkg
2929
Pkg.activate(joinpath(@__DIR__))
30+
using Dates
3031

3132
println("╔═══════════════════════════════════════════════════════════╗")
3233
println("║ ECHIDNA Neural Solver — Training Pipeline ║")
@@ -143,7 +144,9 @@ println("Training ($(training_config.num_epochs) epochs)...")
143144
println("═══════════════════════════════════════════════════════════")
144145

145146
mkpath(save_dir)
147+
t_start = time()
146148
metrics = train_solver!(solver, train_data, val_data; config=training_config)
149+
duration_seconds = time() - t_start
147150

148151
# Save final model
149152
println()
@@ -154,14 +157,93 @@ println("═══════════════════════
154157
final_path = joinpath(save_dir, "final_model")
155158
save_solver(solver, final_path)
156159

160+
# Deploy alias: gnn_endpoint.jl loads gnn_ranker → best_model → final_model.
161+
best_dir = joinpath(save_dir, "best_model")
162+
ranker_dir = joinpath(save_dir, "gnn_ranker")
163+
if isdir(best_dir)
164+
rm(ranker_dir; recursive=true, force=true)
165+
cp(best_dir, ranker_dir)
166+
println("Published best_model → $ranker_dir")
167+
end
168+
157169
# Save vocabulary separately for the API server
158170
BSON.@save joinpath(save_dir, "vocabulary.bson") vocab
159171

172+
# Flat canonical artefacts expected by gnn_endpoint.jl and the spec.
173+
# premise_selector.bson = the full NeuralSolver weights (renamed for clarity).
174+
# tactic_predictor.bson = same weights (tactic side shares the text_encoder).
175+
# vocab.json = human-readable vocabulary for inspection and version checks.
176+
import JSON3 as _JSON3
177+
weights = Flux.state(solver)
178+
BSON.bson(joinpath(save_dir, "premise_selector.bson"), weights=weights)
179+
BSON.bson(joinpath(save_dir, "tactic_predictor.bson"), weights=weights)
180+
181+
# Build a compact vocab.json: token→id map + metadata.
182+
tactic_classes = sort(unique([
183+
ex.proof_state.goal[1:min(8,length(ex.proof_state.goal))]
184+
for ex in vcat(train_data.examples, val_data.examples)
185+
]))
186+
vocab_json = Dict(
187+
"vocab_size" => vocab.vocab_size,
188+
"tactic_classes" => length(tactic_classes),
189+
"token_to_id" => vocab.token_to_id,
190+
)
191+
open(joinpath(save_dir, "vocab.json"), "w") do io
192+
_JSON3.write(io, vocab_json)
193+
end
194+
println("Saved vocab.json ($(vocab.vocab_size) tokens, $(length(tactic_classes)) tactic classes)")
195+
196+
# Compute final MRR / top-k on validation split.
197+
val_metrics = compute_metrics(solver, val_data; k=10)
198+
val_top1 = compute_metrics(solver, val_data; k=1).precision
199+
val_top5 = compute_metrics(solver, val_data; k=5).precision
200+
println("Validation MRR=$(round(val_metrics.mrr, digits=4)) top1=$(round(val_top1, digits=4)) top5=$(round(val_top5, digits=4)) top10=$(round(val_metrics.precision, digits=4))")
201+
202+
# Append a single JSONL row to training_data/metrics_baseline.jsonl.
203+
git_sha = try; strip(read(`git -C $(joinpath(@__DIR__, "..", "..")) rev-parse --short HEAD`, String)); catch; "unknown"; end
204+
metrics_row = Dict{String,Any}(
205+
"timestamp" => string(Dates.now()),
206+
"git_sha" => git_sha,
207+
"mrr" => round(Float64(val_metrics.mrr), digits=6),
208+
"top1" => round(Float64(val_top1), digits=6),
209+
"top5" => round(Float64(val_top5), digits=6),
210+
"top10" => round(Float64(val_metrics.precision), digits=6),
211+
"epochs" => training_config.num_epochs,
212+
"max_proof_states" => max_proof_states,
213+
"duration_seconds" => round(duration_seconds, digits=1),
214+
"device" => has_gpu ? "gpu" : "cpu",
215+
)
216+
metrics_baseline_path = joinpath(data_dir, "metrics_baseline.jsonl")
217+
open(metrics_baseline_path, "a") do io
218+
println(io, _JSON3.write(metrics_row))
219+
end
220+
println("Appended metrics row to $metrics_baseline_path")
221+
222+
# Rewrite models/model_metadata.txt with real values.
223+
metadata_path = joinpath(@__DIR__, "..", "..", "models", "model_metadata.txt")
224+
open(metadata_path, "w") do io
225+
println(io, "# ECHIDNA Neural Models v2.0")
226+
println(io, "# Trained: $(Dates.now())")
227+
println(io, "# Git SHA: $git_sha")
228+
println(io, "# Device: $(has_gpu ? "GPU" : "CPU")")
229+
println(io, "# Premise Selector: vocabulary-based ($(vocab.vocab_size) words)")
230+
println(io, "# Tactic Predictor: neural text encoder ($(length(tactic_classes)) classes)")
231+
println(io, "# MRR: $(round(Float64(val_metrics.mrr), digits=4))")
232+
println(io, "# Top-1 Precision: $(round(Float64(val_top1), digits=4))")
233+
println(io, "# Top-5 Precision: $(round(Float64(val_top5), digits=4))")
234+
println(io, "# Epochs: $(training_config.num_epochs)")
235+
println(io, "# Max Proof States: $max_proof_states")
236+
println(io, "# Training Examples: $(length(train_data.examples))")
237+
println(io, "# Validation Examples: $(length(val_data.examples))")
238+
end
239+
println("Updated $metadata_path")
240+
160241
println()
161242
println("╔═══════════════════════════════════════════════════════════╗")
162243
println("║ Training Complete! ║")
163244
println("╚═══════════════════════════════════════════════════════════╝")
164245
println()
165246
println("Model saved to: $save_dir")
247+
println("MRR: $(round(Float64(val_metrics.mrr), digits=4))")
166248
println("To start the API server:")
167249
println(" julia --project=src/julia src/julia/run_server.jl $save_dir")

src/julia/run_training_cpu.jl

Lines changed: 67 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
using Pkg
2323
Pkg.activate(joinpath(@__DIR__))
24+
using Dates
2425

2526
include(joinpath(@__DIR__, "EchidnaML.jl"))
2627
using .EchidnaML
@@ -80,11 +81,11 @@ config = TrainingConfig(
8081
)
8182

8283
mkpath(save_dir)
83-
train_solver!(solver, train_data, val_data; config=config)
84+
t_start = time()
85+
metrics = train_solver!(solver, train_data, val_data; config=config)
86+
duration_seconds = time() - t_start
8487

8588
# Deploy alias: gnn_endpoint.jl loads gnn_ranker → best_model → final_model.
86-
# Publish the early-stopping best model as the canonical gnn_ranker so the
87-
# server serves the best checkpoint, not the last epoch.
8889
best_dir = joinpath(save_dir, "best_model")
8990
ranker_dir = joinpath(save_dir, "gnn_ranker")
9091
if isdir(best_dir)
@@ -96,4 +97,67 @@ end
9697
save_solver(solver, joinpath(save_dir, "final_model"))
9798
BSON.@save joinpath(save_dir, "vocabulary.bson") vocab
9899

100+
# Flat canonical artefacts expected by gnn_endpoint.jl and the spec.
101+
import JSON3 as _JSON3
102+
weights = Flux.state(solver)
103+
BSON.bson(joinpath(save_dir, "premise_selector.bson"), weights=weights)
104+
BSON.bson(joinpath(save_dir, "tactic_predictor.bson"), weights=weights)
105+
106+
tactic_classes = sort(unique([
107+
ex.proof_state.goal[1:min(8,length(ex.proof_state.goal))]
108+
for ex in vcat(train_data.examples, val_data.examples)
109+
]))
110+
vocab_json = Dict(
111+
"vocab_size" => vocab.vocab_size,
112+
"tactic_classes" => length(tactic_classes),
113+
"token_to_id" => vocab.token_to_id,
114+
)
115+
open(joinpath(save_dir, "vocab.json"), "w") do io
116+
_JSON3.write(io, vocab_json)
117+
end
118+
119+
# Compute final MRR / top-k on validation split.
120+
val_metrics = compute_metrics(solver, val_data; k=10)
121+
val_top1 = compute_metrics(solver, val_data; k=1).precision
122+
val_top5 = compute_metrics(solver, val_data; k=5).precision
123+
println("Validation MRR=$(round(val_metrics.mrr, digits=4)) top1=$(round(val_top1, digits=4)) top5=$(round(val_top5, digits=4)) top10=$(round(val_metrics.precision, digits=4))")
124+
125+
# Append metrics row to training_data/metrics_baseline.jsonl.
126+
git_sha = try; strip(read(`git -C $(joinpath(@__DIR__, "..", "..")) rev-parse --short HEAD`, String)); catch; "unknown"; end
127+
metrics_row = Dict{String,Any}(
128+
"timestamp" => string(Dates.now()),
129+
"git_sha" => git_sha,
130+
"mrr" => round(Float64(val_metrics.mrr), digits=6),
131+
"top1" => round(Float64(val_top1), digits=6),
132+
"top5" => round(Float64(val_top5), digits=6),
133+
"top10" => round(Float64(val_metrics.precision), digits=6),
134+
"epochs" => num_epochs,
135+
"max_proof_states" => max_proof_states,
136+
"duration_seconds" => round(duration_seconds, digits=1),
137+
"device" => "cpu",
138+
)
139+
metrics_baseline_path = joinpath(data_dir, "metrics_baseline.jsonl")
140+
open(metrics_baseline_path, "a") do io
141+
println(io, _JSON3.write(metrics_row))
142+
end
143+
144+
# Rewrite models/model_metadata.txt with real values.
145+
metadata_path = joinpath(@__DIR__, "..", "..", "models", "model_metadata.txt")
146+
open(metadata_path, "w") do io
147+
println(io, "# ECHIDNA Neural Models v2.0")
148+
println(io, "# Trained: $(Dates.now())")
149+
println(io, "# Git SHA: $git_sha")
150+
println(io, "# Device: CPU")
151+
println(io, "# Premise Selector: vocabulary-based ($(vocab.vocab_size) words)")
152+
println(io, "# Tactic Predictor: neural text encoder ($(length(tactic_classes)) classes)")
153+
println(io, "# MRR: $(round(Float64(val_metrics.mrr), digits=4))")
154+
println(io, "# Top-1 Precision: $(round(Float64(val_top1), digits=4))")
155+
println(io, "# Top-5 Precision: $(round(Float64(val_top5), digits=4))")
156+
println(io, "# Epochs: $num_epochs")
157+
println(io, "# Max Proof States: $max_proof_states")
158+
println(io, "# Training Examples: $(length(train_data.examples))")
159+
println(io, "# Validation Examples: $(length(val_data.examples))")
160+
end
161+
99162
println(">>> TRAINING COMPLETE — model saved to $save_dir")
163+
println("MRR: $(round(Float64(val_metrics.mrr), digits=4))")

0 commit comments

Comments
 (0)