tinygrad select model size

31641d10232ed91f87b2279a595227462b1eadc0 · 2024-08-10 12:56:30 +0100 · Alex Cheema

Files touched

M exo/inference/tinygrad/inference.py

Diff

commit 31641d10232ed91f87b2279a595227462b1eadc0
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sat Aug 10 12:56:30 2024 +0100

    tinygrad select model size
---
 exo/inference/tinygrad/inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index d52143fb..ddfebd71 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -92,6 +92,6 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
       return
 
     model_path = await self.shard_downloader.ensure_shard(shard)
-    self.model = build_transformer(model_path, shard, model_size="8B")
+    self.model = build_transformer(model_path, shard, model_size="8B" if "8b" in shard.model_id.lower() else "70B")
     self.tokenizer = AutoTokenizer.from_pretrained(str((model_path if model_path.is_dir() else model_path.parent)))
     self.shard = shard

← e6902b2f add --download-quick-check flag to bypass the hf api calls / · back to Exo · LAPTOP GPU and Laptop GPU prefixes e8267e73 →