zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-12-15 01:42:31 +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

4
llms/mlx_lm/utils.py
View File

@@ -20,7 +20,7 @@ from transformers import PreTrainedTokenizer
# Local imports
from .models.base import KVCache
from .sample_utils import top_p_sampling
from .sample_utils import categorical_sampling, top_p_sampling
from .tokenizer_utils import TokenizerWrapper, load_tokenizer
from .tuner.utils import apply_lora_layers
from .tuner.utils import dequantize as dequantize_model
@@ -169,7 +169,7 @@ def generate_step(
if top_p > 0 and top_p < 1.0:
token = top_p_sampling(logits, top_p, temp)
else:
token = mx.random.categorical(logits * (1 / temp))
token = categorical_sampling(logits, temp)
return token, logprobs