Made temperature and top_p available to the inference engine sample interfaces

52ef6ee4a358f64f61246d8797d15f9edbb26aa2 · 2024-11-11 16:14:47 -0800 · Nel Nibcord

Files touched

M exo/inference/mlx/sharded_inference_engine.py
M exo/inference/tinygrad/inference.py

Diff

commit 52ef6ee4a358f64f61246d8797d15f9edbb26aa2
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Mon Nov 11 16:14:47 2024 -0800

    Made temperature and top_p available to the inference engine sample interfaces
---
 exo/inference/mlx/sharded_inference_engine.py | 4 ++--
 exo/inference/tinygrad/inference.py           | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index e5bd3d4e..dd4e3f44 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -37,10 +37,10 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     self.shard_downloader = shard_downloader
     self.executor = ThreadPoolExecutor(max_workers=1)
 
-  async def sample(self, x) -> np.ndarray:
+  async def sample(self, x, temp: float = 0.0, top_p: float = 1.0) -> np.ndarray:
     y = mx.array(x)
     logits = y[:, -1, :]
-    out = np.array(sample_logits(logits))
+    out = np.array(sample_logits(logits, temp=temp, top_p=top_p))
     return out
 
   async def encode(self, shard: Shard, prompt: str) -> np.ndarray:
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index 9d7fe4ff..14b88505 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -64,10 +64,10 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     self.shard_downloader = shard_downloader
     self.executor = ThreadPoolExecutor(max_workers=1)
 
-  async def sample(self, x: np.ndarray) -> np.ndarray:
+  async def sample(self, x: np.ndarray, temp=TEMPERATURE, top_p: float = 0.0) -> np.ndarray:
     logits = x[:, -1, :]
     def sample_wrapper():
-      return sample_logits(Tensor(logits).flatten(), TEMPERATURE, 0, 0.8, 0.0, 0.0).realize()
+      return sample_logits(Tensor(logits).flatten(), temp, 0, 0.8, top_p, 0.0).realize()
     out = await asyncio.get_running_loop().run_in_executor(self.executor, sample_wrapper)
     return out.numpy()

← 8205a5ae Implemented per-request caching in tinygrad · back to Exo · one-line output buffering 10e9f44a →