| |
|
|
| import json |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
| import torch.nn as nn |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT)) |
|
|
| safetensors = pytest.importorskip("safetensors") |
| from safetensors.torch import load_file, save_file |
|
|
| from train import (checkpoint_speaker_rows, checkpoint_tensor_shapes, load_checkpoint, |
| read_checkpoint) |
|
|
|
|
| def _sd(): |
| g = torch.Generator().manual_seed(0) |
| return {"speaker_encoder.embedding.weight": torch.randn(7, 3, generator=g), |
| "plan_injector.proj.weight": torch.randn(4, 2, generator=g), |
| "lin.weight": torch.randn(3, 3, generator=g), "lin.bias": torch.randn(3, generator=g)} |
|
|
|
|
| def test_read_checkpoint_safetensors_matches_pt(tmp_path): |
| sd = _sd() |
| pt = tmp_path / "c.pt" |
| torch.save({"model_state_dict": sd, "step": 40000, "optimizer_state_dict": {}}, pt) |
| st = tmp_path / "model.safetensors" |
| save_file(sd, str(st), metadata={"step": "40000"}) |
| a, b = read_checkpoint(pt), read_checkpoint(st) |
| assert b["step"] == 40000 and set(b) == {"model_state_dict", "step"} |
| assert set(a["model_state_dict"]) == set(b["model_state_dict"]) |
| for k in sd: |
| assert torch.equal(a["model_state_dict"][k], b["model_state_dict"][k]) |
| assert checkpoint_speaker_rows(pt) == checkpoint_speaker_rows(st) == 7 |
| assert checkpoint_tensor_shapes(pt) == checkpoint_tensor_shapes(st) |
| assert any(k.startswith("plan_injector.") for k in checkpoint_tensor_shapes(st)) |
|
|
|
|
| def test_load_checkpoint_accepts_safetensors(tmp_path): |
| class Tiny(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.lin = nn.Linear(3, 3) |
|
|
| src = Tiny() |
| st = tmp_path / "model.safetensors" |
| save_file({k: v.detach().clone() for k, v in src.state_dict().items()}, str(st), metadata={"step": "7"}) |
| dst = Tiny() |
| assert load_checkpoint(str(st), dst, device="cpu") == 7 |
| for k, v in src.state_dict().items(): |
| assert torch.equal(dst.state_dict()[k], v) |
|
|
|
|
| def _fake_checkpoint(path, **args): |
| t = torch.randn(5, 4) |
| sd = _sd() |
| sd["head.weight"] = t |
| sd["tied.weight"] = t |
| targs = {"hybrid_attention_top_k": 0, "backbone": "mamba", "depth_arch": "mamba2", "depth_layers": 4, |
| "depth_dim": 1024, "depth_cond": "prefix", "depth_feedback": "all"} |
| targs.update(args) |
| torch.save({"model_state_dict": sd, "optimizer_state_dict": {"x": torch.zeros(3)}, "step": 40000, "stage": 2, |
| "args": targs}, path) |
| return sd |
|
|
|
|
| def _export(ck, out, arch): |
| return subprocess.run([sys.executable, str(ROOT / "scripts" / "export_release_weights.py"), "--checkpoint", str(ck), |
| "--arch", arch, "--name", "t", "--out", str(out)], capture_output=True, text=True) |
|
|
|
|
| def test_export_is_bit_exact_and_keeps_aliases(tmp_path): |
| ck = tmp_path / "checkpoint_step_40000.pt" |
| sd = _fake_checkpoint(ck) |
| r = _export(ck, tmp_path / "out", "pure") |
| assert r.returncode == 0, r.stdout + r.stderr |
| back = load_file(str(tmp_path / "out" / "model.safetensors")) |
| assert set(back) == set(sd) |
| for k in sd: |
| assert torch.equal(back[k], sd[k]) |
| cfg = json.loads((tmp_path / "out" / "config.json").read_text()) |
| assert cfg["bit_exact_vs_checkpoint"] and cfg["step"] == 40000 and cfg["speaker_table_rows"] == 7 |
| assert cfg["tied_aliases"] == {"tied.weight": "head.weight"} |
| assert cfg["generator_flags"][-2:] == ["--num_speakers", "7"] and cfg["has_plan_injector"] |
| targs = json.loads((tmp_path / "out" / "training_args.json").read_text()) |
| assert targs["depth_cond"] == "prefix" |
|
|
|
|
| @pytest.mark.parametrize("arch,args", [("hyb", {}), ("pure", {"hybrid_attention_top_k": 4}), |
| ("tfm", {}), ("pure", {"depth_layers": 2})]) |
|
|
|
|
| def test_export_refuses_a_mislabelled_checkpoint(tmp_path, arch, args): |
| ck = tmp_path / "c.pt" |
| _fake_checkpoint(ck, **args) |
| r = _export(ck, tmp_path / "out", arch) |
| assert r.returncode != 0 and "FATAL" in (r.stdout + r.stderr) |
|
|
