pass on tinygrad set_on_download_progress

1d54f105149f909c4720217354b18f46cc4d860b · 2024-07-31 22:50:04 +0100 · Alex Cheema

Files touched

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

Diff

commit 1d54f105149f909c4720217354b18f46cc4d860b
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Jul 31 22:50:04 2024 +0100

    pass on tinygrad set_on_download_progress
---
 exo/inference/inference_engine.py   | 6 +++++-
 exo/inference/tinygrad/inference.py | 5 ++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index 013560a5..827f429b 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from typing import Tuple, Optional
+from typing import Tuple, Optional, Callable
 from abc import ABC, abstractmethod
 from .shard import Shard
 
@@ -13,3 +13,7 @@ class InferenceEngine(ABC):
   @abstractmethod
   async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> Tuple[np.ndarray, str, bool]:
     pass
+
+  @abstractmethod
+  def set_on_download_progress(self, on_download_progress: Callable[[int, int], None]):
+    pass
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index c0e35277..545b335f 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -1,7 +1,7 @@
 import asyncio
 from functools import partial
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Callable
 import json
 import tiktoken
 from tiktoken.load import load_tiktoken_bpe
@@ -294,3 +294,6 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     self.shard = shard
     self.model = model
     self.tokenizer = tokenizer
+
+  def set_on_download_progress(self, on_download_progress: Callable[[int, int], None]):
+    pass
\ No newline at end of file

← d6a7e463 async model downloading with download progress. fixes #102. · back to Exo · fix regression introduced by image_str for tinygrad 76766253 →