Fix tests without pip install (OpenGVLab#45)

Author: carmocca

.github/workflows/cpu-tests.yml

@@ -34,11 +34,13 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
-    - name: Install dependencies
+    - name: Run tests without the package installed
       run: |
-        pip install pytest . -r requirements.txt
+        pip install pytest -r requirements.txt
         pip list
+        pytest --disable-pytest-warnings --strict-markers
 
     - name: Run tests
       run: |
+        pip install . --no-deps
         pytest -v --durations=10 --disable-pytest-warnings --strict-markers

lit_llama/__init__.py

@@ -0,0 +1,2 @@
+from lit_llama.model import LLaMAConfig, LLaMA, RMSNorm, build_rope_cache, apply_rope
+from lit_llama.tokenizer import Tokenizer

scripts/prepare_shakespeare.py

@@ -43,7 +43,7 @@ def prepare(destination_path: Path = Path("data/shakespeare")) -> None:
     train_data = data[: int(n * 0.9)]
     val_data = data[int(n * 0.9) :]
 
-    from lit_llama.tokenizer import Tokenizer
+    from lit_llama import Tokenizer
 
     Tokenizer.train(input=input_file_path, destination=destination_path)
     tokenizer = Tokenizer(destination_path / "tokenizer.model")

tests/conftest.py

@@ -1,13 +1,14 @@
-import os
 import sys
+from pathlib import Path
 
 import pytest
 
+wd = Path(__file__).parent.parent.absolute()
+
 
 @pytest.fixture()
 def orig_llama():
-    wd = os.path.realpath(os.path.join(os.path.dirname(__file__), ".."))
-    sys.path.append(wd)
+    sys.path.append(str(wd))
 
     from scripts.download import download_original
 
@@ -16,3 +17,13 @@ def orig_llama():
     import original_model
 
     return original_model
+
+
+@pytest.fixture()
+def lit_llama():
+    # this adds support for running tests without the package installed
+    sys.path.append(str(wd))
+
+    import lit_llama
+
+    return lit_llama

tests/test_generate.py

@@ -5,7 +5,7 @@
 from io import StringIO
 from pathlib import Path
 from unittest import mock
-from unittest.mock import Mock, PropertyMock, call, ANY
+from unittest.mock import Mock, call, ANY
 
 import pytest
 import torch
@@ -59,13 +59,12 @@ def test_main(tmp_path, monkeypatch):
     tokenizer_path = tmp_path / "tokenizer"
     tokenizer_path.touch()
 
-    class FabricMock(PropertyMock):
+    class FabricMock(Mock):
         @property
         def device(self):
             return torch.device("cpu")
 
-    fabric_mock = FabricMock()
-    monkeypatch.setattr(generate.L, "Fabric", fabric_mock)
+    monkeypatch.setattr(generate.L, "Fabric", FabricMock)
     model_mock = Mock()
     monkeypatch.setattr(generate.LLaMA, "from_name", model_mock)
     load_mock = Mock()
@@ -85,7 +84,6 @@ def device(self):
             checkpoint_path=checkpoint_path,
             tokenizer_path=tokenizer_path,
             model_size="1T",
-            accelerator="litpu",
             temperature=2.0,
             top_k=2,
             num_samples=num_samples,
@@ -96,18 +94,7 @@ def device(self):
     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)
-    model = model_mock.return_value
-    assert fabric_mock.mock_calls == [
-        call(accelerator="litpu", devices=1),
-        call().device.__enter__(),
-        call().device.__exit__(None, None, None),
-        call().setup_module(model),
-    ]
-    model = fabric_mock.return_value.setup_module.return_value
-    assert (
-        generate_mock.mock_calls
-        == [call(model, ANY, 50, model.config.block_size, temperature=2.0, top_k=2)] * num_samples
-    )
+    assert generate_mock.mock_calls == [call(ANY, ANY, 50, ANY, temperature=2.0, top_k=2)] * num_samples
     # only the generated result is printed to stdout
     assert out.getvalue() == "foo bar baz\n" * num_samples
 

tests/test_model.py

@@ -1,7 +1,5 @@
 import torch
 
-import lit_llama.model as lit_llama
-
 
 def copy_mlp(llama_mlp, orig_llama_mlp) -> None:
     orig_llama_mlp.w1.weight.copy_(llama_mlp.c_fc1.weight)
@@ -33,7 +31,7 @@ def copy_weights(llama_model, orig_llama_model) -> None:
 
 
 @torch.no_grad()
-def test_to_orig_llama(orig_llama) -> None:
+def test_to_orig_llama(lit_llama, orig_llama) -> None:
     block_size = 64
     vocab_size = 32000
     n_layer = 16

tests/test_rmsnorm.py

@@ -1,10 +1,8 @@
 import torch
 
-import lit_llama.model as lit_llama
-
 
 @torch.no_grad()
-def test_rmsnorm(orig_llama) -> None:
+def test_rmsnorm(lit_llama, orig_llama) -> None:
     block_size = 16
     vocab_size = 16
 

tests/test_rope.py

@@ -1,7 +1,5 @@
 import torch
 
-import lit_llama.model as lit_llama
-
 
 def build_rope_cache_old(seq_len: int, n_elem: int, dtype: torch.dtype, base: int = 10000) -> torch.Tensor:
     """This is the `build_rope_cache` implementation we initially intended to use, but it is numerically not
@@ -53,7 +51,7 @@ def apply_rope_old(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
 
 
 @torch.no_grad()
-def test_rope(orig_llama) -> None:
+def test_rope(lit_llama, orig_llama) -> None:
     bs, seq_len, n_head, n_embed = 1, 6, 2, 8
     x = torch.randint(0, 10000, size=(bs, seq_len, n_head, n_embed // n_head)).float()