refactor: convert pretrained

Author: dest1n1s

.vscode/extensions.json

@@ -0,0 +1,7 @@
+{
+    "recommendations": [
+        "detachhead.basedpyright",
+        "charliermarsh.ruff",
+        "ms-python.python"
+    ]
+}

pyproject.toml

@@ -21,7 +21,7 @@ build-backend = "hatchling.build"
 
 [dependency-groups]
 dev = [
-    "mypy>=1.13.0",
+    "basedpyright>=1.20.0",
     "pre-commit>=4.0.1",
     "ruff>=0.7.1",
 ]
@@ -131,3 +131,6 @@ docstring-code-format = false
 docstring-code-line-length = "dynamic"
 
 
+[tool.pyright]
+typeCheckingMode = "standard"
+reportUnknownMemberType = false

src/xlens/components/transformer_block.py

@@ -41,10 +41,7 @@ def __init__(self, cfg: HookedTransformerConfig, block_index: int):
         self.layer_id = block_index
 
         if cfg.normalization_type == "LN":
-            normalization_layer: Callable[
-                [HookedTransformerConfig],
-                Callable[[Float[jax.Array, "batch pos d_model"]], Float[jax.Array, "batch pos d_model"]],
-            ] = LayerNorm
+            normalization_layer = LayerNorm
         elif cfg.normalization_type == "LNPre":
             # We've folded in LayerNorm weights, so just need the center + scale parts
             normalization_layer = LayerNormPre

src/xlens/components/unembed.py

@@ -19,7 +19,7 @@ class Unembed(eqx.Module):
     def __init__(self, cfg: HookedTransformerConfig):
         self.cfg = cfg
         # Note that there's a separate variable for d_vocab_out and d_vocab (the input vocab size). For language tasks these are always the same, but for algorithmic tasks we may want them to be different.
-        self.W_U: Float[jax.Array, "d_model d_vocab_out"] = jnp.zeros((self.cfg.d_model, self.cfg.d_vocab_out))
+        self.W_U = jnp.zeros((self.cfg.d_model, self.cfg.d_vocab_out))
 
     def __call__(self, residual: Float[jax.Array, "batch pos d_model"]) -> Float[jax.Array, "batch pos d_vocab_out"]:
         return residual @ self.W_U

src/xlens/config.py

@@ -168,7 +168,7 @@ def __post_init__(self):
         if not self.attn_only:
             if self.d_mlp is None:
                 # For some reason everyone hard codes in this hyper-parameter!
-                self.d_mlp: int = self.d_model * 4
+                self.d_mlp = self.d_model * 4
             assert self.act_fn is not None, "act_fn must be specified for non-attn-only models"
             assert self.act_fn in SUPPORTED_ACTIVATIONS, f"act_fn={self.act_fn} must be one of {SUPPORTED_ACTIVATIONS}"
 

src/xlens/hooked_transformer.py

@@ -7,7 +7,7 @@
 
 from xlens.components import Embed, LayerNorm, LayerNormPre, PosEmbed, RMSNorm, RMSNormPre, TransformerBlock, Unembed
 from xlens.hooks import with_cache, with_hooks
-from xlens.pretrained.loading_from_pretrained import get_pretrained_model_config, get_pretrained_state_dict
+from xlens.pretrained.convert import get_pretrained_model_config, get_pretrained_weights
 from xlens.utils import load_pretrained_weights
 
 from .config import HookedTransformerConfig
@@ -167,7 +167,7 @@ def from_pretrained(cls, model_name: str, hf_model=None) -> "HookedTransformer":
         """
 
         cfg = get_pretrained_model_config(model_name)
-        state_dict = get_pretrained_state_dict(model_name, cfg, hf_model=hf_model)
+        weights = get_pretrained_weights(cfg, model_name, hf_model=hf_model)
         model = HookedTransformer(cfg)
-        model = load_pretrained_weights(model, state_dict)
+        model = load_pretrained_weights(model, weights)
         return model

src/xlens/pretrained/__init__.py

@@ -1,3 +1,3 @@
-from .loading_from_pretrained import get_pretrained_model_config, get_pretrained_state_dict
+from .convert import get_pretrained_model_config, get_pretrained_weights
 
-__all__ = ["get_pretrained_state_dict", "get_pretrained_model_config"]
+__all__ = ["get_pretrained_weights", "get_pretrained_model_config"]

src/xlens/pretrained/convert.py

@@ -0,0 +1,34 @@
+"""Loading Pretrained Models Utilities.
+
+This module contains functions for loading pretrained models from the Hugging Face Hub.
+"""
+
+import jax
+
+from xlens.config import HookedTransformerConfig
+from xlens.pretrained.converters import (
+    GPT2Converter,
+    GPTNeoXConverter,
+    LlamaConverter,
+    MistralConverter,
+    Qwen2Converter,
+)
+from xlens.pretrained.model_converter import HuggingFaceModelConverter
+
+converter = HuggingFaceModelConverter(
+    converters=[
+        GPT2Converter(),
+        Qwen2Converter(),
+        LlamaConverter(),
+        MistralConverter(),
+        GPTNeoXConverter(),
+    ]
+)
+
+
+def get_pretrained_model_config(model_name: str) -> HookedTransformerConfig:
+    return converter.get_pretrained_model_config(model_name)
+
+
+def get_pretrained_weights(cfg: HookedTransformerConfig, model_name: str, hf_model=None) -> dict[str, jax.Array]:
+    return converter.get_pretrained_weights(cfg, model_name, hf_model=hf_model)