fix: parse aliases before get pretrained weight

dest1n1s · Nov 6, 2024 · 3857be8 · 3857be8
1 parent 0cdef98
commit 3857be8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/xlens/pretrained/model_converter.py b/src/xlens/pretrained/model_converter.py
@@ -122,6 +122,9 @@ def _get_safe_weights_files(self, safe_weight_index: Any) -> list[str]:
     def get_pretrained_weights(
         self, cfg: HookedTransformerConfig, model_name_or_path: str, **kwargs: Any
     ) -> dict[str, jax.Array]:
+        model_name_or_path = (
+            model_name_or_path if os.path.isdir(model_name_or_path) else self.rev_alias_map[model_name_or_path]
+        )
         if os.path.isdir(model_name_or_path):
             if os.path.isfile(os.path.join(model_name_or_path, SAFE_WEIGHTS_NAME)):
                 resolved_archive_files = [os.path.join(model_name_or_path, SAFE_WEIGHTS_NAME)]