add missing top_p_sampling import

6659a18e944efb61f4404fbd964a661049613760 · 2024-11-26 15:35:25 +0400 · Alex Cheema

Files touched

M exo/inference/mlx/sharded_inference_engine.py

Diff

commit 6659a18e944efb61f4404fbd964a661049613760
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Tue Nov 26 15:35:25 2024 +0400

    add missing top_p_sampling import
---
 exo/inference/mlx/sharded_inference_engine.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 49c41748..3cf3beaf 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -1,6 +1,7 @@
 import numpy as np
 import mlx.core as mx
 import mlx.nn as nn
+from mlx_lm.sample_utils import top_p_sampling
 from ..inference_engine import InferenceEngine
 from .stateful_model import StatefulModel
 from .sharded_utils import load_shard

← b1b08e68 test add line of code · back to Exo · prio python 3.12 1ab1762e →