giant/tests/test_model_summary.py

"""Tests for `giant model summary` (gitea #46)."""

from __future__ import annotations

from pathlib import Path

import pytest
from typer.testing import CliRunner

from giant import config as gconfig
from giant.cli import app
from giant.materials import MATERIAL_PROPERTIES
from giant.model.summary import _NOT_BUILD_TIME, _built_modules, _vocab_caveats, summarize_model

runner = CliRunner()

_PDG_VOCAB = 300
_MAT_VOCAB = len(MATERIAL_PROPERTIES)


def _cfg(overrides: dict | None = None) -> dict:
    return gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, overrides or {})


@pytest.fixture(scope="module")
def default_summary():
    return summarize_model(_cfg(), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB)


def test_default_config_builds_both_stages_with_a_real_tree(default_summary):
    assert set(default_summary.modules) >= {"stage1", "stage2"}
    for module in default_summary.modules.values():
        assert sum(p.numel() for p in module.parameters()) > 0
    stage1 = default_summary.modules["stage1"]
    assert hasattr(stage1, "cond_enc")
    assert hasattr(stage1, "trunk")
    assert {"input_proj", "blocks", "out_proj"} <= {n for n, _ in stage1.trunk.named_children()}


def test_every_in_scope_leaf_is_classified(default_summary):
    in_scope = {
        p
        for p in gconfig.leaf_paths(gconfig.DEFAULT_CONFIG)
        if p.split(".", 1)[0] in ("conditioning", "stage1_model", "stage2_model")
    }
    classified = set(default_summary.consumed) | set(default_summary.inert) | set(default_summary.elsewhere)
    assert classified == in_scope


def test_not_build_time_allow_list_has_no_stale_entries():
    in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
    stale = set(_NOT_BUILD_TIME) - in_scope
    assert not stale, f"_NOT_BUILD_TIME entries no longer in DEFAULT_CONFIG: {sorted(stale)}"


def test_router_disabled_by_default_so_its_fields_are_inert(default_summary):
    assert "stage1_model.router.n_experts" in default_summary.inert
    assert "stage1_model.router.temperature" in default_summary.inert


def test_markov_history_leaves_attention_dims_inert_but_history_itself_consumed(default_summary):
    assert "stage2_model.autoregressive.attn_n_heads" in default_summary.inert
    assert "stage2_model.autoregressive.attn_n_layers" in default_summary.inert
    assert "stage2_model.autoregressive.history" in default_summary.consumed


def test_single_literal_branch_fields_are_correctly_seen_as_consumed(default_summary):
    """Regression guard: n_sec.owner ("stage2"), n_sec.mode ("head") and
    particle_type.target ("onehot") each branch as `== "one specific other
    literal"` in giant/model/builders.py|models.py. A naive single generic
    sentinel probe lands in the same "not that literal" bucket as the
    current value and never crosses the boundary that actually matters --
    this is exactly what _STRING_ALTERNATIVES exists to fix."""
    assert "stage2_model.n_sec.owner" in default_summary.consumed
    assert "stage2_model.n_sec.mode" in default_summary.consumed
    assert "stage2_model.particle_type.target" in default_summary.consumed


def test_stage1_wgan_generator_swaps_flow_time_dim_for_critic_dims():
    summary = summarize_model(_cfg({"stage1_model": {"generator": "wgan"}}), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB)
    assert "stage1_model.flow.time_dim" in summary.inert
    assert "stage1_model.wgan.noise_dim" in summary.consumed
    assert "stage1_model.wgan.critic_hidden_dim" in summary.consumed


def test_stage2_one_shot_decoder_makes_autoregressive_block_inert():
    summary = summarize_model(
        _cfg({"stage2_model": {"decoder": "one_shot"}}), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB
    )
    assert "stage2_model.autoregressive.history" in summary.inert
    assert "history_encoder" not in {n for n, _ in summary.modules["stage2"].named_children()}


def test_energy_router_enabled_consumes_core_fields_but_not_process_only_fields():
    summary = summarize_model(
        _cfg({"stage1_model": {"router": {"enabled": True, "type": "energy", "n_experts": 4}}}),
        pdg_vocab=_PDG_VOCAB,
        mat_vocab=_MAT_VOCAB,
    )
    assert "stage1_model.router.n_experts" in summary.consumed
    assert "stage1_model.router.temperature" in summary.consumed
    # emb_dim/hidden_dim are pdg/process-router-only kwargs -- build_router's
    # signature filter drops them for an energy router.
    assert "stage1_model.router.hidden_dim" in summary.inert
    assert "stage1_model.router.emb_dim" in summary.inert


def test_vocab_caveat_text_for_embedding_particle_conditioning():
    cfg = _cfg({"conditioning": {"particle": {"type": "embedding"}}})
    caveats = _vocab_caveats(cfg)
    assert any("pdg_vocab" in c and "embedding" in c for c in caveats)
    assert not any("mat_vocab" in c for c in caveats)


def test_pdg_vocab_flag_changes_embedding_table_size():
    cfg = _cfg({"conditioning": {"particle": {"type": "embedding"}}})
    small = _built_modules(cfg, pdg_vocab=10, mat_vocab=_MAT_VOCAB)
    big = _built_modules(cfg, pdg_vocab=1000, mat_vocab=_MAT_VOCAB)
    assert big["stage1"].cond_enc.pdg_emb.weight.numel() > small["stage1"].cond_enc.pdg_emb.weight.numel()


def test_invalid_combo_exits_nonzero_with_validate_config_message(tmp_path: Path):
    config_path = tmp_path / "bad.toml"
    config_path.write_text('[meta]\nconfig_version = 3\n\n[stage2_model.particle_type]\ntarget = "embedding"\n')
    result = runner.invoke(app, ["model", "summary", "--config", str(config_path)])
    assert result.exit_code == 1
    assert "requires conditioning.particle.type = 'embedding'" in result.output


def test_cli_default_smoke():
    result = runner.invoke(app, ["model", "summary"])
    assert result.exit_code == 0, result.output
    assert "stage1" in result.output
    assert "stage2" in result.output
    assert "parameters" in result.output
    assert "trunk" in result.output
    assert "inert under this config" in result.output