Spaces:
Sleeping
Sleeping
File size: 3,044 Bytes
2d52135 | 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 | from pathlib import Path
from unittest import mock
import pytest
from mini_transformer import cli
MINIMAL_CFG = """
model:
name: demo
best_checkpoint_path: ./checkpoints/best.pt
latest_checkpoint_path: null
tokenizer: demo_tok
d_model: 16
num_layers: 1
num_heads: 2
d_ff: 32
dropout_rate: 0.0
vocab_size: 32
max_seq_len: 32
pad_id: 0
bos_id: 1
eos_id: 2
tokenizer:
name: demo_tok
path: ./tokenizer/tokenizer.json
corpus: null
dataset: null
vocab_size: 32
max_seq_len: 32
pad_token: "<pad>"
bos_token: "<bos>"
eos_token: "<eos>"
unk_token: "<unk>"
special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
pad_id: 0
bos_id: 1
eos_id: 2
unk_id: 3
generation:
max_new_tokens: 4
temperature: 1.0
top_k: null
top_p: null
do_sample: false
presence_penalty: 0.0
frequency_penalty: 0.0
no_repeat_ngram: null
min_steps_before_eos: 0
runtime:
seed: 42
device: cpu
output_dir: ./outputs
data_dir: ./data
tokenizer_dir: ./tokenizer
cache_dir: ./cache
checkpoint_path: ./checkpoints
input_text: ""
"""
def _make_demo_model(tmp_path: Path) -> Path:
models_root = tmp_path / "trained_models"
model_dir = models_root / "demo"
config_dir = model_dir / "configs"
config_dir.mkdir(parents=True, exist_ok=True)
(model_dir / "checkpoints").mkdir(exist_ok=True)
(model_dir / "tokenizer").mkdir(exist_ok=True)
(config_dir / "config_inference.yaml").write_text(MINIMAL_CFG)
return models_root
def test_infer_main_uses_local_model(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys):
models_root = _make_demo_model(tmp_path)
monkeypatch.setenv(cli.MODELS_ENV, str(models_root))
with mock.patch.object(cli, "run_inference", return_value=["hello"]) as run_mock:
rc = cli.infer_main(["-m", "demo", "-t", "hi there"])
assert rc == 0
cfg_passed = run_mock.call_args[0][0]
assert Path(cfg_passed.model.best_checkpoint_path).is_absolute()
assert cfg_passed.input_text == "hi there"
out = capsys.readouterr().out
assert "hello" in out
def test_infer_main_defaults_to_first_model(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
models_root = _make_demo_model(tmp_path)
monkeypatch.setenv(cli.MODELS_ENV, str(models_root))
with mock.patch.object(cli, "run_inference", return_value=["out"]) as run_mock:
cli.infer_main(["-t", "text"])
assert run_mock.called
def test_fetch_main_invokes_snapshot(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
monkeypatch.setenv(cli.MODELS_ENV, str(tmp_path / "trained_models"))
with mock.patch("mini_transformer.cli.snapshot_to_local", return_value=tmp_path) as snap:
rc = cli.fetch_main(
[
"AlaBoussoffara/transformer_test",
"--name",
"demo",
"--force",
]
)
assert rc == 0
snap.assert_called_once()
kwargs = snap.call_args.kwargs
assert kwargs["local_name"] == "demo"
assert kwargs["force"] is True
|