fix(ml): make training loop Zygote-differentiable and run

- Replace setdiff-based ranking_loss with BCE loss (differentiable)
- Use Zygote.@ignore for label construction
- Use hcat instead of setindex! for premise vectors (no mutation)
- Pre-encode tokens outside gradient graph
- Rename MultiHeadAttention to EchidnaMHA (avoids Flux name collision)
- Fix NeuralSolver @functor to exclude vocabulary/config from traversal
- Fix TheoremGraph type annotation for SimpleWeightedDiGraph
- Fix compute_metrics to use cosine similarity scoring
- Fix run_training.jl to use module-qualified CUDA/Flux references
- Remove unused Transformers.jl dependency
- Fix empty context/hypothesis handling in encoder

Training runs: Epoch 1 loss 2.20→2.00→2.14 on 53 examples (CPU mode).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: hyperpolymath

src/julia/models/encoder.jl

@@ -533,12 +533,12 @@ function encode_proof_state(state::ProofState, vocab::ProverVocabulary, encoder:
     goal_tokens, goal_features = encode_goal(prover_enc, state.goal)
     goal_ids = encode_tokens(vocab, goal_tokens)
 
-    # Encode context
-    context_tokens = reduce(vcat, [tokenize_text(c) for c in state.context])
+    # Encode context (may be empty)
+    context_tokens = isempty(state.context) ? String[] : reduce(vcat, [tokenize_text(c) for c in state.context])
     context_ids = encode_tokens(vocab, context_tokens)
 
-    # Encode hypotheses
-    hyp_tokens = reduce(vcat, [tokenize_text(h) for h in state.hypotheses])
+    # Encode hypotheses (may be empty)
+    hyp_tokens = isempty(state.hypotheses) ? String[] : reduce(vcat, [tokenize_text(h) for h in state.hypotheses])
     hyp_ids = encode_tokens(vocab, hyp_tokens)
 
     # Combine all tokens with separator

src/julia/models/neural_solver.jl

@@ -39,7 +39,7 @@ Nodes: Theorems, premises, proof states
 Edges: Dependencies, implications, similarity
 """
 struct TheoremGraph
-    graph::SimpleWeightedDiGraph{Int, Float32}
+    graph::SimpleWeightedDiGraph
     node_features::Matrix{Float32}  # (feature_dim, num_nodes)
     node_types::Vector{Symbol}       # :theorem, :premise, :goal
     node_names::Vector{String}
@@ -286,7 +286,7 @@ struct PremiseRanker
     gnn_encoder::GNNEncoder
     goal_encoder::Dense
     premise_encoder::Dense
-    cross_attention::MultiHeadAttention
+    cross_attention  # Our custom EchidnaMHA (not Flux's)
     score_mlp::Chain
 end
 
@@ -298,7 +298,7 @@ function PremiseRanker(feature_dim::Int, hidden_dim::Int, num_gnn_layers::Int=4)
     goal_encoder = Dense(hidden_dim, hidden_dim)
     premise_encoder = Dense(hidden_dim, hidden_dim)
 
-    cross_attention = MultiHeadAttention(hidden_dim, 8)
+    cross_attention = EchidnaMHA(hidden_dim, 8)
 
     score_mlp = Chain(
         Dense(hidden_dim * 2, hidden_dim, relu),
@@ -339,30 +339,30 @@ function (ranker::PremiseRanker)(g::TheoremGraph)
 end
 
 """
-    MultiHeadAttention
+    EchidnaMHA
 
 Standard multi-head attention mechanism.
 """
-struct MultiHeadAttention
+struct EchidnaMHA
     num_heads::Int
     head_dim::Int
     qkv_proj::Dense
     out_proj::Dense
     dropout::Dropout
 end
 
-Flux.@functor MultiHeadAttention
+Flux.@functor EchidnaMHA
 
-function MultiHeadAttention(d_model::Int, num_heads::Int; dropout::Float32=0.1f0)
+function EchidnaMHA(d_model::Int, num_heads::Int; dropout::Float32=0.1f0)
     head_dim = d_model ÷ num_heads
     qkv_proj = Dense(d_model, 3 * d_model)
     out_proj = Dense(d_model, d_model)
     dropout_layer = Dropout(dropout)
 
-    return MultiHeadAttention(num_heads, head_dim, qkv_proj, out_proj, dropout_layer)
+    return EchidnaMHA(num_heads, head_dim, qkv_proj, out_proj, dropout_layer)
 end
 
-function (mha::MultiHeadAttention)(Q::AbstractMatrix, K::AbstractMatrix, V::AbstractMatrix)
+function (mha::EchidnaMHA)(Q::AbstractMatrix, K::AbstractMatrix, V::AbstractMatrix)
     # Q, K, V: (d_model, seq_len)
     seq_len_q = size(Q, 2)
     seq_len_kv = size(K, 2)
@@ -411,7 +411,8 @@ struct NeuralSolver
     config::EchidnaConfig
 end
 
-Flux.@functor NeuralSolver
+# Only traverse neural network fields, not vocabulary/config
+Flux.@functor NeuralSolver (text_encoder, premise_ranker,)
 
 """
     create_solver(vocab::ProverVocabulary; config::EchidnaConfig=get_config())
@@ -529,5 +530,5 @@ end
 
 export TheoremGraph, build_theorem_graph
 export GCNLayer, GraphAttentionLayer, GNNEncoder
-export PremiseRanker, MultiHeadAttention
+export PremiseRanker, EchidnaMHA
 export NeuralSolver, create_solver, save_solver, load_solver

src/julia/run_training.jl

@@ -40,19 +40,20 @@ using .EchidnaML
 
 # Configure for available hardware
 println("Configuring...")
-if CUDA.functional()
+using CUDA, Flux
+has_gpu = CUDA.functional()
+if has_gpu
     println("  GPU: $(CUDA.device())")
-    set_config!(device=gpu)
+    EchidnaML.set_config!(device=Flux.gpu)
 else
     println("  GPU: not available (using CPU)")
-    set_config!(device=cpu)
+    EchidnaML.set_config!(device=Flux.cpu)
 end
 
 # Use smaller model dimensions for CPU training
-config = get_config()
-if config.device === cpu
+if !has_gpu
     println("  Reducing model size for CPU training")
-    set_config!(
+    EchidnaML.set_config!(
         embedding_dim=128,
         hidden_dim=256,
         num_transformer_layers=2,
@@ -61,11 +62,11 @@ if config.device === cpu
     )
 end
 
-println("  Embedding dim: $(get_config().embedding_dim)")
-println("  Hidden dim: $(get_config().hidden_dim)")
-println("  Transformer layers: $(get_config().num_transformer_layers)")
-println("  GNN layers: $(get_config().gnn_num_layers)")
-println("  Batch size: $(get_config().batch_size)")
+println("  Embedding dim: $(EchidnaML.get_config().embedding_dim)")
+println("  Hidden dim: $(EchidnaML.get_config().hidden_dim)")
+println("  Transformer layers: $(EchidnaML.get_config().num_transformer_layers)")
+println("  GNN layers: $(EchidnaML.get_config().gnn_num_layers)")
+println("  Batch size: $(EchidnaML.get_config().batch_size)")
 println()
 
 # Load data