Refactor + tests

src/mechir/modelling/architectures/base/__init__.py

@@ -0,0 +1,11 @@
+from .linear import ClassificationHead, HiddenLinear
+from .components import BertEmbed
+from ._model import HookedEncoder, HookedEncoderForSequenceClassification
+
+__all__ = [
+    "HookedEncoder",
+    "HookedEncoderForSequenceClassification",
+    "ClassificationHead",
+    "HiddenLinear",
+    "BertEmbed",
+]

src/mechir/modelling/architectures/distilbert/__init__.py

@@ -0,0 +1,6 @@
+from ._model import HookedDistilBert, HookedDistilBertForSequenceClassification
+
+__all__ = [
+    "HookedDistilBert",
+    "HookedDistilBertForSequenceClassification",
+]

src/mechir/modelling/architectures/distilbert.py → src/mechir/modelling/architectures/distilbert/_model.py

@@ -12,20 +12,22 @@
 from einops import repeat
 from jaxtyping import Float, Int
 from torch import nn
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, DistilBertModel, DistilBertForSequenceClassification
 from typing_extensions import Literal
 
 from transformer_lens.components import BertBlock, BertMLMHead, Unembed
 from transformer_lens.hook_points import HookPoint
-from mechir.modelling.hooked.components import BertEmbed
-from mechir.modelling.hooked.linear import MLPClassificationHead
-from mechir.modelling.architectures.base import HookedEncoder
+from mechir.modelling.architectures.base.components import BertEmbed
+from mechir.modelling.architectures.base.linear import MLPClassificationHead
+from mechir.modelling.architectures.base._model import HookedEncoder
 from mechir.modelling.hooked.config import HookedTransformerConfig
 
 
-HookedDistilBert = HookedEncoder
+class HookedDistilBert(HookedEncoder):
+    _hf_class = DistilBertModel
 
 class HookedDistilBertForSequenceClassification(HookedDistilBert):
+    _hf_class = DistilBertForSequenceClassification
     """
     This class implements a BERT-style encoder for sequence classification using the components in ./components.py, with HookPoints on every interesting activation. It inherits from HookedDistilBert.
 

src/mechir/modelling/architectures/electra/__init__.py

@@ -0,0 +1,6 @@
+from ._model import HookedElectra, HookedElectraForSequenceClassification
+
+__all__ = [
+    "HookedElectra",
+    "HookedElectraForSequenceClassification",
+]

src/mechir/modelling/architectures/electra.py → src/mechir/modelling/architectures/electra/_model.py

@@ -9,12 +9,13 @@
 import logging
 from typing import Dict, Optional, Union
 
+from transformers import ElectraModel, ElectraForSequenceClassification
 import torch
 from jaxtyping import Float, Int
 from torch import nn
 from transformer_lens.hook_points import HookPoint
-from mechir.modelling.hooked.linear import ClassificationHead, HiddenLinear
-from mechir.modelling.architectures.base import HookedEncoder
+from mechir.modelling.architectures.base.linear import ClassificationHead, HiddenLinear
+from mechir.modelling.architectures.base._model import HookedEncoder
 from mechir.modelling.hooked.config import HookedTransformerConfig
 
 
@@ -38,7 +39,10 @@ def forward(self, resid: Float[torch.Tensor, "batch d_model"]) -> torch.Tensor:
         post_act = self.hook_post(self.activation(pre_act))
         return self.out_proj(post_act)
 
-HookedElectra = HookedEncoder
+
+class HookedElectra(HookedEncoder):
+    _hf_class = ElectraModel
+
 
 class HookedElectraForSequenceClassification(HookedEncoder):
     """
@@ -49,6 +53,7 @@ class HookedElectraForSequenceClassification(HookedEncoder):
         - There is no preprocessing (e.g. LayerNorm folding) when loading a pretrained model
         - The model only accepts tokens as inputs, and not strings, or lists of strings
     """
+    _hf_class = ElectraForSequenceClassification
 
     def __init__(self, cfg, tokenizer=None, move_to_device=True, **kwargs):
         super().__init__(cfg, tokenizer, move_to_device, **kwargs)