From 2f443cc6d73ffadf2484d4367345fcf415eed350 Mon Sep 17 00:00:00 2001
From: Prince Canuma <prince.gdt@gmail.com>
Date: Sat, 14 Dec 2024 16:39:46 +0100
Subject: [PATCH] add tests

---
 llms/tests/test_models.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/llms/tests/test_models.py b/llms/tests/test_models.py
index 374a5113..d6decb3f 100644
--- a/llms/tests/test_models.py
+++ b/llms/tests/test_models.py
@@ -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()