| 9 | |
| 10 | class TestMistral(unittest.TestCase): |
| 11 | def test_model(self): |
| 12 | vocab_size = 100 |
| 13 | L = 32 |
| 14 | args = mistral.ModelArgs( |
| 15 | dim=128, |
| 16 | n_layers=2, |
| 17 | head_dim=32, |
| 18 | hidden_dim=256, |
| 19 | n_heads=4, |
| 20 | n_kv_heads=4, |
| 21 | norm_eps=1e-3, |
| 22 | vocab_size=vocab_size, |
| 23 | ) |
| 24 | |
| 25 | model = mistral.Mistral(args) |
| 26 | inputs = mx.random.randint(0, vocab_size, (L,)) |
| 27 | logits, cache = model(inputs[None]) |
| 28 | self.assertEqual(logits.shape, [1, L, vocab_size]) |
| 29 | self.assertEqual(logits.dtype, mx.float32) |
| 30 | self.assertEqual(len(cache), args.n_layers) |
| 31 | |
| 32 | params = tree_map(lambda p: p.astype(mx.float16), model.parameters()) |
| 33 | model.update(params) |
| 34 | logits, _ = model(inputs[None]) |
| 35 | self.assertEqual(logits.dtype, mx.float16) |
| 36 | |
| 37 | def test_generate(self): |
| 38 | model, tokenizer = mistral.load_model("mistral-7B-v0.1") |