fix: tokenize

0730532351b99167f4352a34968423da9ed8af51 · 2024-10-11 03:02:51 +0530 · Varshith

Files touched

M exo/inference/mlx/sharded_inference_engine.py

Diff

commit 0730532351b99167f4352a34968423da9ed8af51
Author: Varshith <varshith.bathini@sprinklr.com>
Date:   Fri Oct 11 03:02:51 2024 +0530

    fix: tokenize
---
 exo/inference/mlx/sharded_inference_engine.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 7b920ccc..122628d5 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -8,6 +8,7 @@ from typing import Optional
 from exo.download.shard_download import ShardDownloader
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 
 class MLXDynamicShardInferenceEngine(InferenceEngine):
   def __init__(self, shard_downloader: ShardDownloader):
@@ -20,7 +21,8 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     loop = asyncio.get_running_loop()
     if image_str:
       image = await get_image_from_str(image_str)
-      inputs = await loop.run_in_executor(self.executor, self.tokenizer, prompt, image, return_tensors="np")
+      tokenize = partial(self.tokenizer, prompt, image, return_tensors="np")
+      inputs = await loop.run_in_executor(self.executor, tokenize)
       pixel_values = mx.array(inputs["pixel_values"])
       input_ids = mx.array(inputs["input_ids"])
       output_data: np.ndarray = np.array(await loop.run_in_executor(self.executor, self.stateful_sharded_model.step, request_id, input_ids, pixel_values))

← 17065d87 dynamic halfway partition point in unit test · back to Exo · chore: Support multiple nodes download progress section in t 53ea5767 →