chore(mlx-lm): fix the top_p implementation.

[mzbac]

The top_p is supposed to sort prob in ascending order and select the cumulative token prob above the threshold. The current implementation works because we have a bug in reversing the sorted prob, which accidentally makes sorted prob in the correct order. So, it's better to fix the implementation to reduce the confusion.

[awni]

Btw @mzbac looks like the CI issue is resolved for you, thanks!!

[mzbac]

@awni I have added the unit test, so it is ready for review again. Please let me know if there is anything else that needs to be updated.

[awni]

I'm not sure about introducing a separate file for that. Do you think it's necessary? Is the idea that we will add more samplers?

[mzbac]

The main reason for putting it into a separate file is the current utils.py has become too large, making it difficult to maintain. So, I think we may need to start splitting up the utils.py file. sampling utils might be a good starting point. Sampling utils would include top_p, repetitive penalty and any future sampling methods that we want to support.

[awni]

So, I think we may need to start splitting up the utils.py file.

Agreed.

Sampling utils would include top_p, repetitive penalty and any future sampling methods that we want to support.

I'm not wild about the name sampling_utils but I think collecting that related functionality in a new file makes sense. Let's go with it for now but I may change the file name in the future if I can think of a better one 😄

[awni]

🙏 thanks for fixing that and for the tests!