mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 09:21:18 +08:00
fbed720d6f
* chore(mlx-lm): clean up the top p imp * chore: clean up * chore: add test * chore: address comments * chore: clean up docs string * chore: clean up test
38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import mlx.core as mx
|
|
from mlx_lm.sample_utils import top_p_sampling
|
|
|
|
|
|
class TestLora(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
|
|
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(mx.allclose(args[0], mx.log(expected_top_probs)))
|
|
|
|
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, 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(mx.allclose(args[0], mx.log(expected_top_probs)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|