Since infer_prompt is a thin wrapper that works the same for all inference engines, we can de-abstract it

b9d0fb68258a0bfd7e76f3b83f3924167524f796 · 2024-11-11 00:39:55 -0800 · Nel Nibcord

Files touched

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

Diff

commit b9d0fb68258a0bfd7e76f3b83f3924167524f796
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Mon Nov 11 00:39:55 2024 -0800

    Since infer_prompt is a thin wrapper that works the same for all inference engines, we can de-abstract it
---
 exo/inference/inference_engine.py             | 10 ++++++----
 exo/inference/mlx/sharded_inference_engine.py |  5 -----
 exo/inference/tinygrad/inference.py           |  5 -----
 3 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index cb93c7a5..4349889c 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -20,13 +20,15 @@ class InferenceEngine(ABC):
   async def decode(self, shard: Shard, tokens: np.ndarray) -> str:
     pass
 
-  @abstractmethod
-  async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> np.ndarray:
-    pass
-
   @abstractmethod
   async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
     pass
+  
+  async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> np.ndarray:
+    tokens = await self.encode(shard, prompt)
+    output_data = await self.infer_tensor(request_id, shard, tokens, inference_state)
+    return output_data 
+
 
 def get_inference_engine(inference_engine_name: str, shard_downloader: 'ShardDownloader'):
   if DEBUG >= 2:
diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 98eae461..10f0cd73 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -53,11 +53,6 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.decode, tokens)
     return tokens
     
-  async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> (np.ndarray, bool):
-    tokens = await self.encode(shard, prompt)
-    output_data = await self.infer_tensor(request_id, shard, tokens, inference_state)
-    return output_data 
-
   async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> (np.ndarray, bool):
     await self.ensure_shard(shard)
     output_data: np.ndarray = np.array(await asyncio.get_running_loop().run_in_executor(self.executor, self.model, mx.array(input_data), request_id))
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index eef1386c..3834ab1e 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -82,11 +82,6 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.decode, tokens)
     return tokens
 
-  async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> np.ndarray:
-    tokens = await self.encode(shard, prompt)
-    output_data = await self.infer_tensor(request_id, shard, tokens, inference_state)
-    return output_data 
-
   async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> tuple[np.ndarray, str, bool]:
     await self.ensure_shard(shard)
     start_pos = json.loads(inference_state or "{}").get("start_pos", 0)

← 527c7a6e Applied new interface to tinygrad and dummy inference engine · back to Exo · Make sure they're np arrays 9b66758b →