mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-08-30 02:36:37 +08:00
add tests
This commit is contained in:
parent
52595dafae
commit
2f443cc6d7
@ -851,6 +851,19 @@ class TestModels(unittest.TestCase):
|
|||||||
model = exaone.Model(args)
|
model = exaone.Model(args)
|
||||||
self.model_test_runner(model, args.model_type, args.vocab_size, args.num_layers)
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
Loading…
Reference in New Issue
Block a user