Преглед изворни кода

tinygrad select model size

Alex Cheema пре 1 година
родитељ
комит
31641d1023
1 измењених фајлова са 1 додато и 1 уклоњено
  1. 1 1
      exo/inference/tinygrad/inference.py

+ 1 - 1
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