load safetensors first

InternLM · Sep 26, 2023 · b640544 · b640544
1 parent cb82fe3
commit b640544
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/lmdeploy/serve/turbomind/deploy.py b/lmdeploy/serve/turbomind/deploy.py
@@ -117,10 +117,15 @@ def load_checkpoint(model_path):
     Returns:
         Dict[str, torch.Tensor]: weight in torch format
     """
-    files = [
-        file for file in os.listdir(model_path)
-        if file.endswith('.bin') or file.endswith('.safetensors')
-    ]
+    suffixes = ['.safetensors', '.bin']
+    for suffix in suffixes:
+        files = [
+            file for file in os.listdir(model_path) if file.endswith(suffix)
+        ]
+        if len(files) > 0:
+            break
+
+    assert len(files) > 0, f'could not find checkpoints in {model_path}'
     files = sorted(files)
     print(files)
     params = {}