Skip to content

Commit d2c64fd

Browse files
hyperpolymathclaude
hyperpolymath
and
claude
committed
feat(ml): wire Julia ML pipeline — data loading, training, and serving
Complete the neurosymbolic training pipeline: - training/dataloader.jl: JSONL → TrainingExample pipeline that reads proof_states, premises, and tactics JSONL files, joins by proof_id, samples negative premises, builds vocabulary, and produces batched TrainingDataset for the Flux training loop - run_training.jl: standalone entry point that loads data, creates NeuralSolver (GNN + Transformer), trains with ranking + contrastive loss, saves model + vocabulary to disk - run_server.jl: standalone entry point that loads a trained model and serves the REST API on port 8090 (matching what server.rs expects) - Project.toml: add 9 missing dependencies (Flux, CUDA, GNN, etc.) and SimpleWeightedGraphs, Dates - train.jl: fix gradient clipping (use Optimisers.ClipNorm chain), replace CSV metrics with JSONL, remove stub load_training_data - Fix SPDX headers to PMPL-1.0-or-later on all Julia files Usage: julia --project=src/julia src/julia/run_training.jl [data_dir] [save_dir] julia --project=src/julia src/julia/run_server.jl [model_dir] [port] Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 2d8c07f commit d2c64fd

10 files changed

Lines changed: 629 additions & 51 deletions

File tree

‎src/julia/EchidnaML.jl‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# SPDX-FileCopyrightText: 2025 ECHIDNA Project Team
2-
# SPDX-License-Identifier: MIT AND Palimpsest-0.6
2+
# SPDX-License-Identifier: PMPL-1.0-or-later
33

44
"""
55
EchidnaML
@@ -185,6 +185,7 @@ get_config() = CONFIG[]
185185
# Include submodules
186186
include("models/encoder.jl")
187187
include("models/neural_solver.jl")
188+
include("training/dataloader.jl")
188189
include("training/train.jl")
189190
include("inference/predict.jl")
190191
include("api/server.jl")

‎src/julia/Project.toml‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,33 @@
1+
# SPDX-License-Identifier: PMPL-1.0-or-later
2+
# EchidnaML — Neural premise selection for 30 theorem provers
3+
14
[deps]
5+
BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
6+
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
7+
DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
8+
Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
9+
Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
10+
GraphNeuralNetworks = "cffab07f-9bc2-4a1a-8404-210e54f16cc5"
11+
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
212
HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
313
JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
414
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
515
Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
16+
NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
17+
Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
18+
Oxygen = "df9a0d86-3283-4920-82dc-4555fc0d1d8b"
19+
ProgressMeter = "92933f4c-e287-5a05-a399-4b506bd553cd"
20+
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
21+
Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
22+
SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
23+
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
24+
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
625

726
[compat]
827
julia = "1.10"
28+
Flux = "0.16"
929
HTTP = "1"
1030
JSON3 = "1"
31+
CUDA = "5"
32+
Optimisers = "0.4"
33+
Oxygen = "1"

‎src/julia/api/server.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# SPDX-FileCopyrightText: 2025 ECHIDNA Project Team
2-
# SPDX-License-Identifier: MIT AND Palimpsest-0.6
2+
# SPDX-License-Identifier: PMPL-1.0-or-later
33

44
"""
55
api/server.jl

‎src/julia/inference/predict.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# SPDX-FileCopyrightText: 2025 ECHIDNA Project Team
2-
# SPDX-License-Identifier: MIT AND Palimpsest-0.6
2+
# SPDX-License-Identifier: PMPL-1.0-or-later
33

44
"""
55
inference/predict.jl

‎src/julia/models/encoder.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# SPDX-FileCopyrightText: 2025 ECHIDNA Project Team
2-
# SPDX-License-Identifier: MIT AND Palimpsest-0.6
2+
# SPDX-License-Identifier: PMPL-1.0-or-later
33

44
"""
55
models/encoder.jl

‎src/julia/models/neural_solver.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# SPDX-FileCopyrightText: 2025 ECHIDNA Project Team
2-
# SPDX-License-Identifier: MIT AND Palimpsest-0.6
2+
# SPDX-License-Identifier: PMPL-1.0-or-later
33

44
"""
55
models/neural_solver.jl

‎src/julia/run_server.jl‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# SPDX-License-Identifier: PMPL-1.0-or-later
2+
# Copyright (c) 2026 Jonathan D.A. Jewell (hyperpolymath) <j.d.a.jewell@open.ac.uk>
3+
#
4+
# run_server.jl — Start the ECHIDNA neural premise selection API server.
5+
#
6+
# Usage:
7+
# julia --project=src/julia src/julia/run_server.jl [model_dir] [port]
8+
#
9+
# Defaults:
10+
# model_dir = models/neural/final_model
11+
# port = 8090
12+
#
13+
# The Rust server (server.rs) connects to this on ECHIDNA_ML_API_URL
14+
# (default http://127.0.0.1:8090).
15+
16+
using Pkg
17+
Pkg.activate(joinpath(@__DIR__))
18+
19+
println("╔═══════════════════════════════════════════════════════════╗")
20+
println("║ ECHIDNA Neural Solver — API Server ║")
21+
println("╚═══════════════════════════════════════════════════════════╝")
22+
println()
23+
24+
# Parse arguments
25+
model_dir = length(ARGS) >= 1 ? ARGS[1] : joinpath(@__DIR__, "..", "..", "models", "neural", "final_model")
26+
port = length(ARGS) >= 2 ? parse(Int, ARGS[2]) : 8090
27+
28+
println("Model directory: $model_dir")
29+
println("Port: $port")
30+
println()
31+
32+
# Load module
33+
println("Loading EchidnaML module...")
34+
include(joinpath(@__DIR__, "EchidnaML.jl"))
35+
using .EchidnaML
36+
37+
# Start server (blocking — runs until Ctrl+C)
38+
start_api_server(model_dir; port=port, host="0.0.0.0", cache_size=1000, async=false)

‎src/julia/run_training.jl‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
# SPDX-License-Identifier: PMPL-1.0-or-later
2+
# Copyright (c) 2026 Jonathan D.A. Jewell (hyperpolymath) <j.d.a.jewell@open.ac.uk>
3+
#
4+
# run_training.jl — Main entry point for training the ECHIDNA neural solver.
5+
#
6+
# Usage:
7+
# julia --project=src/julia src/julia/run_training.jl [data_dir] [save_dir]
8+
#
9+
# Defaults:
10+
# data_dir = training_data/
11+
# save_dir = models/neural/
12+
#
13+
# This script:
14+
# 1. Loads JSONL training data (proof states + premises)
15+
# 2. Builds vocabulary from the corpus
16+
# 3. Creates the NeuralSolver (GNN + Transformer)
17+
# 4. Trains with ranking + contrastive loss
18+
# 5. Saves the trained model + vocabulary to save_dir
19+
20+
using Pkg
21+
Pkg.activate(joinpath(@__DIR__))
22+
23+
println("╔═══════════════════════════════════════════════════════════╗")
24+
println("║ ECHIDNA Neural Solver — Training Pipeline ║")
25+
println("╚═══════════════════════════════════════════════════════════╝")
26+
println()
27+
28+
# Parse arguments
29+
data_dir = length(ARGS) >= 1 ? ARGS[1] : joinpath(@__DIR__, "..", "..", "training_data")
30+
save_dir = length(ARGS) >= 2 ? ARGS[2] : joinpath(@__DIR__, "..", "..", "models", "neural")
31+
32+
println("Data directory: $data_dir")
33+
println("Save directory: $save_dir")
34+
println()
35+
36+
# Load module
37+
println("Loading EchidnaML module...")
38+
include(joinpath(@__DIR__, "EchidnaML.jl"))
39+
using .EchidnaML
40+
41+
# Configure for available hardware
42+
println("Configuring...")
43+
if CUDA.functional()
44+
println(" GPU: $(CUDA.device())")
45+
set_config!(device=gpu)
46+
else
47+
println(" GPU: not available (using CPU)")
48+
set_config!(device=cpu)
49+
end
50+
51+
# Use smaller model dimensions for CPU training
52+
config = get_config()
53+
if config.device === cpu
54+
println(" Reducing model size for CPU training")
55+
set_config!(
56+
embedding_dim=128,
57+
hidden_dim=256,
58+
num_transformer_layers=2,
59+
gnn_num_layers=2,
60+
batch_size=16
61+
)
62+
end
63+
64+
println(" Embedding dim: $(get_config().embedding_dim)")
65+
println(" Hidden dim: $(get_config().hidden_dim)")
66+
println(" Transformer layers: $(get_config().num_transformer_layers)")
67+
println(" GNN layers: $(get_config().gnn_num_layers)")
68+
println(" Batch size: $(get_config().batch_size)")
69+
println()
70+
71+
# Load data
72+
println("═══════════════════════════════════════════════════════════")
73+
println("Loading training data...")
74+
println("═══════════════════════════════════════════════════════════")
75+
76+
train_data, val_data, vocab = load_training_data(data_dir;
77+
train_split=0.8f0,
78+
max_proof_states=50000, # Cap for reasonable training time
79+
num_negatives=20
80+
)
81+
82+
if isempty(train_data.examples)
83+
println("ERROR: No training data loaded. Check that JSONL files exist in $data_dir")
84+
exit(1)
85+
end
86+
87+
println()
88+
println("Vocabulary size: $(vocab.vocab_size)")
89+
println("Training examples: $(length(train_data.examples))")
90+
println("Validation examples: $(length(val_data.examples))")
91+
println()
92+
93+
# Create solver
94+
println("═══════════════════════════════════════════════════════════")
95+
println("Creating NeuralSolver...")
96+
println("═══════════════════════════════════════════════════════════")
97+
98+
solver = create_solver(vocab)
99+
println("Model created successfully")
100+
println()
101+
102+
# Configure training
103+
training_config = TrainingConfig(
104+
num_epochs=30,
105+
learning_rate=1f-4,
106+
lr_schedule=:cosine,
107+
weight_decay=1f-5,
108+
gradient_clip_norm=1.0f0,
109+
loss_alpha=0.5f0,
110+
early_stopping_patience=5,
111+
checkpoint_every=5,
112+
eval_every=1,
113+
save_dir=save_dir
114+
)
115+
116+
# Train
117+
println("═══════════════════════════════════════════════════════════")
118+
println("Training ($(training_config.num_epochs) epochs)...")
119+
println("═══════════════════════════════════════════════════════════")
120+
121+
mkpath(save_dir)
122+
metrics = train_solver!(solver, train_data, val_data; config=training_config)
123+
124+
# Save final model
125+
println()
126+
println("═══════════════════════════════════════════════════════════")
127+
println("Saving final model...")
128+
println("═══════════════════════════════════════════════════════════")
129+
130+
final_path = joinpath(save_dir, "final_model")
131+
save_solver(solver, final_path)
132+
133+
# Save vocabulary separately for the API server
134+
BSON.@save joinpath(save_dir, "vocabulary.bson") vocab
135+
136+
println()
137+
println("╔═══════════════════════════════════════════════════════════╗")
138+
println("║ Training Complete! ║")
139+
println("╚═══════════════════════════════════════════════════════════╝")
140+
println()
141+
println("Model saved to: $save_dir")
142+
println("To start the API server:")
143+
println(" julia --project=src/julia src/julia/run_server.jl $save_dir")

0 commit comments

Comments
 (0)