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

fix(mlx-lm): sorted probs in top_p implementation. (#610)

Browse Source 
* fix(mlx-lm): the top p imp

* chore: address comment
This commit is contained in:
Anchen
2024-03-26 09:07:55 +11:00
committed by GitHub
parent bbfcc103d7
commit 0ab01b4626
2 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
llms/tests/test_sample_utils.py
View File

@@ -18,6 +18,7 @@ class TestLora(unittest.TestCase):
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)))
logits = mx.array([[1.0, 2.0, 3.0, 4.0]])
@@ -30,6 +31,7 @@ class TestLora(unittest.TestCase):
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)))