zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-10-28 17:08: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/mlx_lm/sample_utils.py
View File

@@ -22,7 +22,7 @@ def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.arr
# sort probs in ascending order
sorted_indices = mx.argsort(probs, axis=-1)
sorted_probs = probs[..., sorted_indices]
sorted_probs = probs[..., sorted_indices.squeeze(0)]
cumulative_probs = mx.cumsum(sorted_probs, axis=-1)