| | import functools |
| | import subprocess |
| | import sys |
| | from contextlib import contextmanager, redirect_stdout |
| | from io import StringIO |
| | from pathlib import Path |
| | from unittest import mock |
| | from unittest.mock import Mock, call, ANY |
| |
|
| | import torch |
| |
|
| | wd = Path(__file__).parent.parent.absolute() |
| |
|
| |
|
| | @functools.lru_cache(maxsize=1) |
| | def load_generate_script(): |
| | sys.path.append(str(wd)) |
| |
|
| | import generate as generate |
| |
|
| | return generate |
| |
|
| |
|
| | def test_generate(): |
| | generate = load_generate_script() |
| |
|
| | from lit_llama.model import LLaMA, LLaMAConfig |
| |
|
| | T, C = 5, 3 |
| | logits = torch.randn(T, C) |
| | input_idx = torch.randint(10, size=(T,)) |
| |
|
| | config = LLaMAConfig(block_size=128, vocab_size=16, n_layer=1, n_head=4, n_embd=8) |
| | model = LLaMA(config) |
| | max_new_tokens = 20 |
| |
|
| | multinomial_results = [] |
| | original_multinomial = torch.multinomial |
| |
|
| | def multinomial(*args, **kwargs): |
| | out = original_multinomial(*args, **kwargs) |
| | multinomial_results.append(out) |
| | return out |
| |
|
| | with mock.patch("torch.multinomial", multinomial): |
| | out = generate.generate(model, input_idx, max_new_tokens, max_seq_length=10, top_k=4) |
| |
|
| | assert out.size(0) == T + max_new_tokens |
| | multinomial_results = torch.hstack(multinomial_results) |
| | expected = torch.cat((input_idx, multinomial_results)) |
| | assert out.shape == expected.shape |
| | torch.testing.assert_close(out, expected) |
| |
|
| |
|
| | @mock.patch("torch.cuda.is_bf16_supported", return_value=False) |
| | def test_main(tmp_path, monkeypatch): |
| | generate = load_generate_script() |
| |
|
| | checkpoint_path = tmp_path / "ckpt" |
| | checkpoint_path.touch() |
| | tokenizer_path = tmp_path / "tokenizer" |
| | tokenizer_path.touch() |
| |
|
| | class FabricMock(Mock): |
| | @property |
| | def device(self): |
| | return torch.device("cpu") |
| |
|
| | @contextmanager |
| | def init_module(self, empty_init): |
| | yield |
| |
|
| | monkeypatch.setattr(generate.L, "Fabric", FabricMock) |
| | model_mock = Mock() |
| | monkeypatch.setattr(generate.LLaMA, "from_name", model_mock) |
| | lookup_mock = Mock(return_value="1T") |
| | monkeypatch.setattr(generate, "llama_model_lookup", lookup_mock) |
| | load_mock = Mock() |
| | load_mock.return_value = load_mock |
| | load_mock.__enter__ = Mock() |
| | load_mock.__exit__ = Mock() |
| | monkeypatch.setattr(generate.torch, "load", load_mock) |
| | monkeypatch.setattr(generate, "lazy_load", load_mock) |
| | tokenizer_mock = Mock() |
| | tokenizer_mock.return_value.encode.return_value = torch.tensor([[1, 2, 3]]) |
| | tokenizer_mock.return_value.decode.return_value = "foo bar baz" |
| | monkeypatch.setattr(generate, "Tokenizer", tokenizer_mock) |
| | generate_mock = Mock() |
| | generate_mock.return_value = torch.tensor([[3, 2, 1]]) |
| | monkeypatch.setattr(generate, "generate", generate_mock) |
| |
|
| | num_samples = 2 |
| | out = StringIO() |
| | with redirect_stdout(out): |
| | generate.main( |
| | checkpoint_path=checkpoint_path, |
| | tokenizer_path=tokenizer_path, |
| | temperature=2.0, |
| | top_k=2, |
| | num_samples=num_samples, |
| | ) |
| |
|
| | model_mock.assert_called_once_with("1T") |
| | load_mock.assert_called_once_with(checkpoint_path) |
| | tokenizer_mock.assert_called_once_with(tokenizer_path) |
| | assert len(tokenizer_mock.return_value.decode.mock_calls) == num_samples |
| | assert torch.allclose(tokenizer_mock.return_value.decode.call_args[0][0], generate_mock.return_value) |
| | assert generate_mock.mock_calls == [call(ANY, ANY, 50, temperature=2.0, top_k=2)] * num_samples |
| | |
| | assert out.getvalue() == "foo bar baz\n" * num_samples |
| |
|
| |
|
| | def test_cli(): |
| | cli_path = wd / "generate.py" |
| | output = subprocess.check_output([sys.executable, cli_path, "-h"]) |
| | output = str(output.decode()) |
| | assert "Generates text samples" in output |
| |
|
