Corrected type annotations

c06b5f3b56c4271764ce2769ea49e579150c4f17 · 2024-11-11 00:48:17 -0800 · Nel Nibcord
Files touched

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

commit c06b5f3b56c4271764ce2769ea49e579150c4f17
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Mon Nov 11 00:48:17 2024 -0800

    Corrected type annotations
---
 exo/inference/dummy_inference_engine.py       |  6 +-----
 exo/inference/mlx/sharded_inference_engine.py | 10 +++++-----
 exo/inference/tinygrad/inference.py           |  8 ++++----
 3 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/exo/inference/dummy_inference_engine.py b/exo/inference/dummy_inference_engine.py
index 74d20513..be0cc94e 100644
--- a/exo/inference/dummy_inference_engine.py
+++ b/exo/inference/dummy_inference_engine.py
@@ -28,11 +28,7 @@ class DummyInferenceEngine(InferenceEngine):
   async def decode(self, shard: Shard, tokens: np.ndarray) -> str:
     return ' '.join([random_string(np.random.randint(1, 34)) for token in tokens])
 
-  async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None):
-    output_data = await self.infer_tensor(request_id, shard, await self.encode(shard, prompt), 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]:
+  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
     await self.ensure_shard(shard)
     sequence_length = input_data.shape[0 if self.shard.is_first_layer() else 1]
     output = np.random.random(size=(1, sequence_length, self.vocab_size if self.shard.is_last_layer() else self.hidden_size))
diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 10f0cd73..e5bd3d4e 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -37,23 +37,23 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     self.shard_downloader = shard_downloader
     self.executor = ThreadPoolExecutor(max_workers=1)
 
-  async def sample(self, x):
+  async def sample(self, x) -> np.ndarray:
     y = mx.array(x)
     logits = y[:, -1, :]
     out = np.array(sample_logits(logits))
     return out
 
-  async def encode(self, shard: Shard, prompt: str):
+  async def encode(self, shard: Shard, prompt: str) -> np.ndarray:
     await self.ensure_shard(shard)
     tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.encode, prompt)
-    return tokens
+    return np.array(tokens)
 
-  async def decode(self, shard: Shard, tokens):
+  async def decode(self, shard: Shard, tokens) -> str:
     await self.ensure_shard(shard)
     tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.decode, tokens)
     return tokens
     
-  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> (np.ndarray, bool):
+  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
     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))
     return output_data
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index cd12eb91..3725a82b 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -65,24 +65,24 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     self.shard_downloader = shard_downloader
     self.executor = ThreadPoolExecutor(max_workers=1)
 
-  async def sample(self, x: np.ndarray):
+  async def sample(self, x: np.ndarray) -> np.ndarray:
     logits = x[:, -1, :]
     def sample_wrapper():
       return sample_logits(Tensor(x).flatten(), TEMPERATURE, 0, 0.8, 0.0, 0.0).realize()
     out = await asyncio.get_running_loop().run_in_executor(self.executor, sample_wrapper)
     return out.numpy()
 
-  async def encode(self, shard: Shard, prompt: str):
+  async def encode(self, shard: Shard, prompt: str) -> np.ndarray:
     await self.ensure_shard(shard)
     tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.encode, prompt)
     return np.array(tokens)
   
-  async def decode(self, shard: Shard, tokens):
+  async def decode(self, shard: Shard, tokens) -> str:
     await self.ensure_shard(shard)
     tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.decode, tokens)
     return tokens
 
-  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> tuple[np.ndarray, str, bool]:
+  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
     await self.ensure_shard(shard)
     start_pos = json.loads(inference_state or "{}").get("start_pos", 0)
     output_data = await asyncio.get_running_loop().run_in_executor(self.executor, self.model, Tensor(input_data), start_pos)
← 9b66758b Make sure they're np arrays · back to Exo · I think this is more faithful to how it was originally done aefc0d7c →