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

Faster sampling with mx.compile (#937)

Browse Source 
* faster sampling with compile

* fix test
This commit is contained in:
Awni Hannun
2024-08-15 11:29:09 -07:00
committed by GitHub
parent 95840f32e2
commit 9b83004631
3 changed files with 32 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
llms/tests/test_sample_utils.py
View File

@@ -1,38 +1,32 @@
import unittest
from unittest.mock import patch
import mlx.core as mx
from mlx_lm.sample_utils import top_p_sampling
class TestSamplingUtils(unittest.TestCase):
@patch("mlx.core.random.categorical")
def test_top_p_sampling(self, mock_categorical):
logits = mx.array([[1.0, 2.0, 3.0, 4.0]])
top_p = 0.3
def test_top_p_sampling(self):
probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
logits = mx.log(probs)
temperature = 1.0
expected_token = mx.array([3])
mock_categorical.return_value = expected_token
token = top_p_sampling(logits, top_p, temperature)
expected_top_probs = mx.array([[0.0, 0.0, 0.0, 0.643914]])
self.assertTrue(mx.allclose(token, expected_token))
args, _ = mock_categorical.call_args
self.assertTrue(args[0].shape == expected_top_probs.shape)
self.assertTrue(mx.allclose(args[0], mx.log(expected_top_probs)))
token = top_p_sampling(logits, 0.3, temperature).item()
self.assertEqual(token, 0)
logits = mx.array([[1.0, 2.0, 3.0, 4.0]])
top_p = 0.9
temperature = 1.0
expected_token = mx.array([3])
mock_categorical.return_value = expected_token
token = top_p_sampling(logits, 0.95, temperature).item()
self.assertTrue(token in (0, 3))
token = top_p_sampling(logits, top_p, temperature)
expected_top_probs = mx.array([[0.0, 0.0871443, 0.236883, 0.643914]])
self.assertTrue(mx.allclose(token, expected_token))
args, _ = mock_categorical.call_args
self.assertTrue(args[0].shape == expected_top_probs.shape)
self.assertTrue(mx.allclose(args[0], mx.log(expected_top_probs)))
probs = mx.array([0.0, 0.5, 0.4, 0.1])[None]
logits = mx.log(probs)
token = top_p_sampling(logits, 0.4, temperature).item()
self.assertEqual(token, 1)
token = top_p_sampling(logits, 0.6, temperature).item()
self.assertTrue(token in (1, 2))
token = top_p_sampling(logits, 0.95, temperature).item()
self.assertTrue(token in (1, 2, 3))
if __name__ == "__main__":