diff --git a/contracts/contracts.nt b/contracts/contracts.nt index 218e0a585d..c870b20a71 100644 --- a/contracts/contracts.nt +++ b/contracts/contracts.nt @@ -13998,6 +13998,11 @@ "tracing-observability-v1"^^ . "tracing-observability"^^ . "1.0.0"^^ . + . + . + "contracts/train-loader-try-load-apr-v1.yaml"^^ . + "train-loader-try-load-apr-v1"^^ . + "pattern"^^ . . . . diff --git a/contracts/shapes.ttl b/contracts/shapes.ttl index a718856fa0..bed403e57d 100644 --- a/contracts/shapes.ttl +++ b/contracts/shapes.ttl @@ -1983,11 +1983,11 @@ ] ; sh:property [ sh:path ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:in ( "8.9" "12.1" ) ; @@ -2067,11 +2067,11 @@ sh:maxCount 1 ; sh:class ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2084,11 +2084,11 @@ sh:targetClass ; sh:property [ sh:path ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2103,11 +2103,11 @@ sh:targetClass ; sh:property [ sh:path ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2194,7 +2194,7 @@ sh:minCount 1 ; sh:maxCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; sh:property [ sh:path ; @@ -2217,7 +2217,7 @@ ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2239,7 +2239,7 @@ sh:minCount 1 ; sh:maxCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; sh:property [ sh:path ; @@ -2257,7 +2257,7 @@ sh:path ; sh:minCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; sh:property [ sh:path ; @@ -2269,7 +2269,7 @@ ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2283,7 +2283,7 @@ ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2347,11 +2347,11 @@ sh:path ; sh:minCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2392,11 +2392,11 @@ sh:path ; sh:minCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2448,11 +2448,11 @@ sh:path ; sh:minCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2511,11 +2511,11 @@ sh:minCount 1 ; sh:maxCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2543,11 +2543,11 @@ sh:minCount 1 ; sh:maxCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2569,11 +2569,11 @@ sh:minCount 1 ; sh:maxCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2595,11 +2595,11 @@ sh:minCount 1 ; sh:maxCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2626,11 +2626,11 @@ sh:minCount 1 ; sh:maxCount 1 ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2721,7 +2721,7 @@ sh:minCount 1 ; sh:class ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; sh:property [ sh:path ; @@ -2733,7 +2733,7 @@ ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; @@ -2768,11 +2768,11 @@ sh:maxCount 1 ; sh:class ; sh:nodeKind sh:IRI ; - sh:node ; + sh:node ; ] ; . - a sh:NodeShape ; + a sh:NodeShape ; sh:property [ sh:path ; sh:minCount 1 ; diff --git a/contracts/train-loader-try-load-apr-v1.yaml b/contracts/train-loader-try-load-apr-v1.yaml new file mode 100644 index 0000000000..5d2025f91c --- /dev/null +++ b/contracts/train-loader-try-load-apr-v1.yaml @@ -0,0 +1,36 @@ +metadata: + version: 1.0.0 + created: '2026-10-18' + author: PAIML Engineering + description: "config::train::loader try_load_apr — loading a transformer from an APR file or checkpoint directory: unreadable or absent input is None, the checkpoint step comes from the checkpoint_step metadata and otherwise from the model-step-N file name, a directory resolves to its newest step, and the weights round-trip." + kind: pattern + references: + - "crates/aprender-train/src/config/train/loader/data.rs — try_load_apr, try_load_apr_for_inference, find_latest_apr_checkpoint" + +equations: + try_load_apr_step_and_weights: + formula: "step == metadata.checkpoint_step if present else parse(file name model-step-N.apr) else 0; a directory picks max N; absent, corrupt or empty-dir input yields None" + domain: "a path (APR file or directory) and a TransformerConfig" + codomain: "Option<(Option, usize)>" + invariants: + - "Metadata step takes precedence over the file-name step" + - "Loaded embedding weights equal the saved ones" + +proof_obligations: + - id: TLA-INV-001 + type: invariant + property: "Metadata checkpoint_step wins over the file-name step" + formal: "∀ t ∈ {test_try_load_apr_reads_step_from_metadata_over_filename}: exit(t) = 0" + applies_to: try_load_apr_step_and_weights + - id: TLA-INV-002 + type: invariant + property: "A directory resolves to its newest model-step-N.apr" + formal: "∀ t ∈ {test_try_load_apr_directory_picks_latest_step_from_filename}: exit(t) = 0" + applies_to: try_load_apr_step_and_weights + +falsification_tests: + - id: FALSIFY-TLA-001 + rule: "metadata key ignored" + prediction: "reading a different metadata key makes the step fall back to the file name and turns the metadata test RED" + test: "cargo test -p aprender-train --lib test_try_load_apr_reads_step_from_metadata_over_filename" + if_fails: "the checkpoint_step metadata path is unguarded" diff --git a/crates/aprender-train/src/config/train/loader/tests.rs b/crates/aprender-train/src/config/train/loader/tests.rs index f67d7eea2b..16bb6d52a1 100644 --- a/crates/aprender-train/src/config/train/loader/tests.rs +++ b/crates/aprender-train/src/config/train/loader/tests.rs @@ -2096,3 +2096,78 @@ fn test_tabular_shapes_name_the_offending_batch_index() { let err = validate_tabular_batch_shapes(&batches, 8).expect_err("batch 2 must be rejected"); assert!(err.to_string().contains("batch 2"), "error must name the batch: {err}"); } + +// ── try_load_apr (via try_load_apr_for_inference) ──────────────────────── + +/// Write `model` as an APR file at `path` with the given extra metadata. +fn write_tiny_apr(path: &std::path::Path, model: &Transformer, meta: &[(&str, &str)]) { + use aprender::serialization::apr::AprWriter; + let mut w = AprWriter::new(); + for (k, v) in meta { + w.set_metadata(k.to_string(), serde_json::json!(v)); + } + for (name, t) in model.named_parameters() { + let data = t.data(); + w.add_tensor_f32(name, vec![data.len()], data.as_slice().expect("contiguous")); + } + w.write(path).expect("write apr"); +} + +#[test] +fn test_try_load_apr_rejects_missing_path_and_corrupt_file() { + let cfg = TransformerConfig::tiny(); + let dir = tempfile::tempdir().expect("tempdir"); + assert!(try_load_apr_for_inference(&dir.path().join("absent.apr"), &cfg).is_none()); + let bad = dir.path().join("bad.apr"); + std::fs::write(&bad, b"not an apr file").expect("write"); + assert!(try_load_apr_for_inference(&bad, &cfg).is_none()); + // Empty directory: no checkpoint to pick. + assert!(try_load_apr_for_inference(dir.path(), &cfg).is_none()); +} + +#[test] +fn test_try_load_apr_reads_step_from_metadata_over_filename() { + let cfg = TransformerConfig::tiny(); + let model = Transformer::new(&cfg); + let dir = tempfile::tempdir().expect("tempdir"); + let p = dir.path().join("model-step-5.apr"); + write_tiny_apr(&p, &model, &[("checkpoint_step", "42")]); + let (loaded, step) = try_load_apr_for_inference(&p, &cfg).expect("loads"); + assert_eq!(step, 42, "metadata step wins over filename step"); + assert_eq!(loaded.expect("model").named_parameters().len(), model.named_parameters().len()); +} + +#[test] +fn test_try_load_apr_directory_picks_latest_step_from_filename() { + let cfg = TransformerConfig::tiny(); + let model = Transformer::new(&cfg); + let dir = tempfile::tempdir().expect("tempdir"); + write_tiny_apr(&dir.path().join("model-step-3.apr"), &model, &[]); + write_tiny_apr(&dir.path().join("model-step-12.apr"), &model, &[]); + let (_, step) = try_load_apr_for_inference(dir.path(), &cfg).expect("loads"); + assert_eq!(step, 12, "no metadata: step parsed from the newest file name"); +} + +#[test] +fn test_try_load_apr_round_trips_weights_with_training_tensor_present() { + let cfg = TransformerConfig::tiny(); + let model = Transformer::new(&cfg); + let dir = tempfile::tempdir().expect("tempdir"); + let p = dir.path().join("model.apr"); + { + use aprender::serialization::apr::AprWriter; + let mut w = AprWriter::new(); + for (name, t) in model.named_parameters() { + let d = t.data(); + w.add_tensor_f32(name, vec![d.len()], d.as_slice().expect("contiguous")); + } + w.add_tensor_f32("__training__.opt_state", vec![2], &[1.0, 2.0]); + w.write(&p).expect("write"); + } + let (loaded, step) = try_load_apr_for_inference(&p, &cfg).expect("loads"); + assert_eq!(step, 0, "model.apr carries no step"); + let loaded = loaded.expect("model"); + let a = model.embed_tokens.weight.data(); + let b = loaded.embed_tokens.weight.data(); + assert_eq!(a.as_slice(), b.as_slice(), "embedding weights round-trip"); +}