Correctly setting safetensor path if the stanford model is a local fo…

…lder (#68)
lightonai · Oct 22, 2024 · b34bba0 · b34bba0
1 parent b89517f
commit b34bba0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pylate/models/Dense.py b/pylate/models/Dense.py
@@ -121,6 +121,9 @@ def from_stanford_weights(
                 token=token,
                 use_auth_token=use_auth_token,
             )
+        # If the model a local folder, load the safetensor
+        else:
+            model_name_or_path = os.path.join(model_name_or_path, "model.safetensors")
         with safe_open(model_name_or_path, framework="pt", device="cpu") as f:
             state_dict = {"linear.weight": f.get_tensor("linear.weight")}