lars
da7cde3ef9
CI / Lint (ruff check) (push) Successful in 27s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 36s
CI / Type check (ty) (push) Successful in 39s
CI / Format (ruff format) (pull_request) Successful in 30s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 30s
CI / Tests (pull_request) Successful in 2m50s
CI / Tests (push) Successful in 2m58s
Works through docs/v0.3.0-followups.md item by item, closing the gap between the design doc and the shipped v0.3.0-stage2-autoregressive code: 1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2 dispatch, stage-2 particle-type-class marginal. 2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run. 3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/ train instead of the hardcoded K_MAX constant. 4. Mixed conditioning.particle.type / conditioning.material.type support end-to-end (data pipeline + dwarf warm-cache). 5. conditioning.share_stages = true: one shared ConditionEncoder instance across both stages. 6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2 (was silently unimplemented). 7. giant predict/rollout: implement conditioning.*.type = "onehot" via the checkpoint's saved pdg_topn_map/mat_topn_map. 8. network.py's checkpoint-path model_config migration now fails loudly on non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's TOML-load path (§4.2). 9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a rollout-capable checkpoint (§9). Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics), mostly a test-helper dict-unpack pattern that made every unrelated constructor keyword look like a type error. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md
|
|
decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must
|
|
independently override each stage's config block, and must take precedence
|
|
over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
|
|
apply the same value to both stages for backward compatibility."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from typer.testing import CliRunner
|
|
|
|
import giant.cli as cli
|
|
|
|
runner = CliRunner()
|
|
|
|
|
|
def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dict:
|
|
captured: dict = {}
|
|
|
|
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
|
captured["cfg"] = cfg
|
|
|
|
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
|
|
|
result = runner.invoke(
|
|
cli.app,
|
|
["train", "dummy.parquet", "--out", str(tmp_path / "run")] + args,
|
|
)
|
|
assert result.exit_code == 0, result.output
|
|
return captured["cfg"]
|
|
|
|
|
|
def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):
|
|
cfg = _invoke_and_capture_cfg(
|
|
monkeypatch,
|
|
tmp_path,
|
|
["--mode", "wgan", "--stage1-generator", "flow"],
|
|
)
|
|
assert cfg["stage1_model"]["generator"] == "flow"
|
|
assert cfg["stage2_model"]["generator"] == "wgan"
|
|
|
|
|
|
def test_stage2_only_knobs(monkeypatch, tmp_path):
|
|
cfg = _invoke_and_capture_cfg(
|
|
monkeypatch,
|
|
tmp_path,
|
|
[
|
|
"--stage2-decoder",
|
|
"one_shot",
|
|
"--stage2-k-max",
|
|
"8",
|
|
"--stage2-hidden-dim",
|
|
"32",
|
|
"--stage2-context-dim",
|
|
"16",
|
|
"--stage2-stage1-context",
|
|
"sampled",
|
|
],
|
|
)
|
|
assert cfg["stage2_model"]["decoder"] == "one_shot"
|
|
assert cfg["stage2_model"]["k_max"] == 8
|
|
assert cfg["stage2_model"]["hidden_dim"] == 32
|
|
assert cfg["stage2_model"]["context_dim"] == 16
|
|
assert cfg["stage2_model"]["stage1_context"] == "sampled"
|
|
# untouched stage1 defaults
|
|
assert cfg["stage1_model"]["hidden_dim"] == 256
|
|
|
|
|
|
def test_stage1_hidden_dim_flag_overrides_legacy_hidden_dim_flag(monkeypatch, tmp_path):
|
|
cfg = _invoke_and_capture_cfg(
|
|
monkeypatch,
|
|
tmp_path,
|
|
["--hidden-dim", "64", "--stage1-hidden-dim", "128"],
|
|
)
|
|
assert cfg["stage1_model"]["hidden_dim"] == 128
|
|
|
|
|
|
def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
|
|
cfg = _invoke_and_capture_cfg(
|
|
monkeypatch,
|
|
tmp_path,
|
|
[
|
|
"--mode",
|
|
"wgan",
|
|
"--n-critic",
|
|
"5",
|
|
"--stage1-n-critic",
|
|
"3",
|
|
"--stage2-gp-weight",
|
|
"2.5",
|
|
],
|
|
)
|
|
assert cfg["stage1_model"]["wgan"]["n_critic"] == 3
|
|
assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
|
|
assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
|
|
assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
|