zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-28 23:49:43 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add tests

Browse Source
This commit is contained in:
Prince Canuma 2024-12-14 16:39:46 +01:00
parent 52595dafae
commit 2f443cc6d7
1 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
llms/tests/test_models.py
View File

@ -851,6 +851,19 @@ class TestModels(unittest.TestCase):
model = exaone.Model(args)
self.model_test_runner(model, args.model_type, args.vocab_size, args.num_layers)
def test_cohere2(self):
from mlx_lm.models import cohere2
args = cohere2.ModelArgs(
model_type="cohere2",
hidden_size=4096,
head_dim=128,
num_hidden_layers=40,
sliding_window=4096,
sliding_window_pattern=4,
)
model = cohere2.Model(args)
self.model_test_runner(model, args.model_type, args.vocab_size, args.num_hidden_layers)
if __name__ == "__main__":
unittest.main()