Spaces:
Sleeping
Sleeping
File size: 7,572 Bytes
595cac0 5313f7c 595cac0 489a9c3 5313f7c 489a9c3 5313f7c 595cac0 266b0a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | import pytest
from plastic.cli import build_parser, main
from plastic.store import ArtifactStore
def test_help_lists_commands(capsys):
with pytest.raises(SystemExit):
build_parser().parse_args(["--help"])
out = capsys.readouterr().out
for cmd in ("data", "train", "models", "calibrate", "chat"):
assert cmd in out
def test_calibrate_dispatches_pretrained_chat_models_to_the_real_chat_calibration(tmp_path, capsys, monkeypatch):
from plastic.harness import calibrate as calibrate_mod
from plastic.harness.calibration_prompts import DEFAULT_CALIBRATION_PROMPTS
root = str(tmp_path / "artifacts")
ArtifactStore(root).register_model("chat_m", {"backend": "ttt", "domain": "text"})
calls = []
class FakeCal:
n_chunks = 3
thresholds = {"chunk_loss": 1.0}
def fake_qwen(store, model_id, prompts, **kw):
calls.append((model_id, list(prompts), kw))
return FakeCal()
monkeypatch.setattr(calibrate_mod, "calibrate_qwen", fake_qwen)
monkeypatch.setattr(calibrate_mod, "calibrate_model", lambda *a, **k: pytest.fail("toy path used for a pretrained record"))
# no --prompts: the bundled benign set
assert main(["calibrate", "chat_m", "--artifacts-root", root, "--fpr", "0.2", "--max-new-tokens", "8"]) == 0
out = capsys.readouterr()
assert "bundled" in out.err and '"chunk_loss"' in out.out
assert calls[-1][0] == "chat_m" and calls[-1][1] == list(DEFAULT_CALIBRATION_PROMPTS)
assert calls[-1][2]["target_fpr"] == 0.2 and calls[-1][2]["max_new_tokens"] == 8 and calls[-1][2]["cusum_prompts"] is None
# explicit prompt files
(tmp_path / "p.txt").write_text("alpha\nbeta\n", encoding="utf-8")
(tmp_path / "c.json").write_text('["gamma"]', encoding="utf-8")
assert main(["calibrate", "chat_m", "--artifacts-root", root, "--prompts", str(tmp_path / "p.txt"), "--cusum-prompts", str(tmp_path / "c.json")]) == 0
assert calls[-1][1] == ["alpha", "beta"] and calls[-1][2]["cusum_prompts"] == ["gamma"]
def test_default_calibration_prompts_are_distinct_nonempty_strings():
from plastic.harness.calibration_prompts import DEFAULT_CALIBRATION_PROMPTS
assert len(DEFAULT_CALIBRATION_PROMPTS) >= 32
assert len(set(DEFAULT_CALIBRATION_PROMPTS)) == len(DEFAULT_CALIBRATION_PROMPTS)
assert all(isinstance(p, str) and p.strip() == p and len(p) > 10 for p in DEFAULT_CALIBRATION_PROMPTS)
def test_read_prompts_accepts_lines_or_json_list(tmp_path):
from plastic.cli import _read_prompts
lines = tmp_path / "p.txt"
lines.write_text("first prompt\n\n second prompt \n", encoding="utf-8")
assert _read_prompts(str(lines)) == ["first prompt", "second prompt"]
js = tmp_path / "p.json"
js.write_text('["a", "b"]', encoding="utf-8")
assert _read_prompts(str(js)) == ["a", "b"]
bad = tmp_path / "bad.json"
bad.write_text('[1, 2]', encoding="utf-8")
with pytest.raises(SystemExit):
_read_prompts(str(bad))
def test_train_physics_via_cli(tmp_path, capsys):
root = str(tmp_path / "artifacts")
rc = main(
[
"train", "physics",
"--artifacts-root", root,
"--steps", "3", "--batch-size", "2", "--seq-len", "32", "--episodes-per-seq", "2",
"--d-model", "32", "--heads", "2", "--layers", "1", "--chunk", "16",
"--eval-every", "0", "--save-every", "0", "--eval-batches", "1", "--log-every", "1",
"--device", "cpu", "--no-muon", "--warmup-steps", "1",
]
)
assert rc == 0
out = capsys.readouterr().out.strip().splitlines()
model_id = out[-1]
assert model_id.startswith("phys_")
assert ArtifactStore(root).load_model_record(model_id)["status"] == "completed"
rc = main(["models", "--artifacts-root", root])
assert rc == 0
assert model_id in capsys.readouterr().out
def test_calibrate_session_chat_and_physics_via_cli(tmp_path, capsys):
import os
from plastic.data.text import encode_documents_to_bin
from plastic.tokenizer.bpe import Tokenizer
root = str(tmp_path / "artifacts")
d = str(tmp_path / "data")
os.makedirs(d)
docs = ["alpha beta gamma delta epsilon " * 80, "one two three four five six " * 80]
tok = Tokenizer.train(docs, vocab_size=300)
tok.save(os.path.join(d, "tokenizer.json"))
encode_documents_to_bin(tok, docs, os.path.join(d, "train.bin"))
encode_documents_to_bin(tok, docs, os.path.join(d, "validation.bin"))
assert main(["train", "text", "--data", d, "--artifacts-root", root, "--model-id", "lm_t", "--steps", "2",
"--batch-size", "2", "--seq-len", "32", "--d-model", "32", "--heads", "2", "--layers", "1", "--chunk", "8",
"--eval-every", "0", "--save-every", "0", "--eval-batches", "1", "--log-every", "1", "--device", "cpu",
"--warmup-steps", "1", "--mqar-frac", "0"]) == 0
capsys.readouterr()
assert main(["calibrate", "lm_t", "--artifacts-root", root, "--data", d, "--chunks", "24", "--fisher-chunks", "4", "--fpr", "0.1"]) == 0
out = capsys.readouterr().out
assert "thresholds" in out and os.path.exists(os.path.join(root, "models", "lm_t", "calibration.json"))
assert main(["session", "new", "--model", "lm_t", "--session-id", "s1", "--artifacts-root", root]) == 0
assert capsys.readouterr().out.strip() == "s1"
assert main(["chat", "s1", "alpha beta gamma delta epsilon alpha", "--artifacts-root", root, "--max-new-tokens", "4", "--seed", "0"]) == 0
capsys.readouterr()
assert main(["session", "fork", "s1", "s2", "--artifacts-root", root]) == 0
assert capsys.readouterr().out.strip() == "s2"
assert main(["session", "list", "--artifacts-root", root]) == 0
out = capsys.readouterr().out
assert "s1" in out and "s2" in out
assert main(["session", "show", "s1", "--artifacts-root", root]) == 0
assert "n_transactions" in capsys.readouterr().out
# physics
assert main(["train", "physics", "--artifacts-root", root, "--model-id", "ph_t", "--steps", "2", "--batch-size", "2",
"--seq-len", "32", "--episodes-per-seq", "2", "--d-model", "32", "--heads", "2", "--layers", "1", "--chunk", "8",
"--eval-every", "0", "--save-every", "0", "--eval-batches", "1", "--log-every", "1", "--device", "cpu",
"--no-muon", "--warmup-steps", "1"]) == 0
capsys.readouterr()
assert main(["calibrate", "ph_t", "--artifacts-root", root, "--chunks", "16", "--fisher-chunks", "2", "--fpr", "0.1"]) == 0
capsys.readouterr()
assert main(["session", "new", "--model", "ph_t", "--session-id", "p1", "--artifacts-root", root]) == 0
capsys.readouterr()
assert main(["physics", "p1", "--artifacts-root", root, "--steps", "24", "--mu", "0.1"]) == 0
out = capsys.readouterr().out
assert "adaptive_mse" in out
def test_sleep_cli_pins_the_replay_revision_and_accepts_an_override():
"""ASTRA-181: --replay-revision defaults to the pinned SmolTalk commit; an empty string follows the Hub's main."""
from plastic.cli import build_parser
from plastic.sleep import SMOLTALK_REVISION
p = build_parser()
ns = p.parse_args(["sleep", "m"])
assert ns.replay_revision == SMOLTALK_REVISION
ns = p.parse_args(["sleep", "m", "--replay-revision", "abc1234", "--flagged-policy", "downweight", "--flagged-weight", "0.5"])
assert (ns.replay_revision, ns.flagged_policy, ns.flagged_weight) == ("abc1234", "downweight", 0.5)
assert p.parse_args(["sleep", "m", "--replay-revision", ""]).replay_revision == ""
|