Dynamically add boot function to bridge (#964)

* Dynamically add boot function to bridge

* Fix imports

* Fix circular import

* Keep mypy happy

* Edit tests to use TransformerBridge.boot_transformers instead of boot

* Remove exposing transformers in model_bridge init

* removed pretrained and replaced with nn module as param

* removed extra import and ran import

* added lazy import again

---------

Co-authored-by: Bryce Meyer <bryce13950@gmail.com>

Author: degenfabian

tests/integration/model_bridge/test_bridge_integration.py

@@ -8,13 +8,13 @@
 import torch
 
 from transformer_lens.ActivationCache import ActivationCache
-from transformer_lens.boot import boot
+from transformer_lens.model_bridge import TransformerBridge
 
 
 def test_model_initialization():
     """Test that the model can be initialized correctly."""
     model_name = "gpt2"  # Use a smaller model for testing
-    bridge = boot(model_name)
+    bridge = TransformerBridge.boot_transformers(model_name)
 
     assert bridge is not None, "Bridge should be initialized"
     assert bridge.tokenizer is not None, "Tokenizer should be initialized"
@@ -24,7 +24,7 @@ def test_model_initialization():
 def test_text_generation():
     """Test basic text generation functionality."""
     model_name = "gpt2"  # Use a smaller model for testing
-    bridge = boot(model_name)
+    bridge = TransformerBridge.boot_transformers(model_name)
 
     if bridge.tokenizer.pad_token is None:
         bridge.tokenizer.pad_token = bridge.tokenizer.eos_token
@@ -39,7 +39,7 @@ def test_text_generation():
 def test_hooks():
     """Test that hooks can be added and removed correctly."""
     model_name = "gpt2"  # Use a smaller model for testing
-    bridge = boot(model_name)
+    bridge = TransformerBridge.boot_transformers(model_name)
 
     if bridge.tokenizer.pad_token is None:
         bridge.tokenizer.pad_token = bridge.tokenizer.eos_token
@@ -77,7 +77,7 @@ def test_hook(tensor, hook):
 def test_cache():
     """Test that the cache functionality works correctly."""
     model_name = "gpt2"  # Use a smaller model for testing
-    bridge = boot(model_name)
+    bridge = TransformerBridge.boot_transformers(model_name)
 
     if bridge.tokenizer.pad_token is None:
         bridge.tokenizer.pad_token = bridge.tokenizer.eos_token
@@ -105,7 +105,7 @@ def test_cache():
 def test_component_access():
     """Test that model components can be accessed correctly."""
     model_name = "gpt2"  # Use a smaller model for testing
-    bridge = boot(model_name)
+    bridge = TransformerBridge.boot_transformers(model_name)
 
     # Test accessing various components
     assert hasattr(bridge, "embed"), "Bridge should have embed component"

transformer_lens/__init__.py

@@ -18,7 +18,6 @@
 from .BertNextSentencePrediction import BertNextSentencePrediction
 from . import head_detector
 from . import loading_from_pretrained as loading
-from . import boot
 from . import patching
 from . import train
 
@@ -41,5 +40,4 @@
     "EasyTransformerConfig",
     "EasyTransformerKeyValueCache",
     "EasyTransformerKeyValueCacheEntry",
-    "boot",
 ]

transformer_lens/model_bridge/__init__.py

@@ -3,12 +3,10 @@
 This module provides functionality to bridge between different model architectures.
 """
 
-from transformer_lens.factories.architecture_adapter_factory import (
-    ArchitectureAdapterFactory,
-)
 from transformer_lens.model_bridge.architecture_adapter import (
     ArchitectureAdapter,
 )
+
 from transformer_lens.model_bridge.bridge import (
     TransformerBridge,
 )
@@ -39,9 +37,11 @@
     TransformerLensPath,
 )
 
+import transformer_lens.model_bridge.sources.transformers
+
+
 __all__ = [
     "ArchitectureAdapter",
-    "ArchitectureAdapterFactory",
     "TransformerBridge",
     "AttentionBridge",
     "BlockBridge",

transformer_lens/model_bridge/architecture_adapter.py

@@ -6,7 +6,7 @@
 from typing import Any, cast
 
 import torch
-from transformers.modeling_utils import PreTrainedModel
+from torch import nn
 
 from transformer_lens.model_bridge.conversion_utils.conversion_steps import (
     WeightConversionSet,
@@ -307,7 +307,7 @@ def translate_transformer_lens_path(
             return remote_path.split(".")[-1]
         return remote_path
 
-    def convert_weights(self, hf_model: PreTrainedModel) -> dict[str, torch.Tensor]:
+    def convert_weights(self, hf_model: nn.Module) -> dict[str, torch.Tensor]:
         """Convert the weights from the HuggingFace format to the HookedTransformer format.
 
         Args:

transformer_lens/model_bridge/sources/__init__.py

@@ -0,0 +1,4 @@
+"""Sources module.
+
+This module provides functionality to load and convert models from HuggingFace to TransformerLens format.
+"""

transformer_lens/boot.py …ens/model_bridge/sources/transformers.pytransformer_lens/boot.py renamed to transformer_lens/model_bridge/sources/transformers.py transformer_lens/boot.py renamed to transformer_lens/model_bridge/sources/transformers.py

@@ -1,4 +1,4 @@
-"""Boot module for TransformerLens.
+"""Transformers module for TransformerLens.
 
 This module provides functionality to load and convert models from HuggingFace to TransformerLens format.
 """
@@ -14,7 +14,6 @@
     PreTrainedTokenizerBase,
 )
 
-from transformer_lens.model_bridge import ArchitectureAdapterFactory
 from transformer_lens.model_bridge.bridge import TransformerBridge
 from transformer_lens.utils import get_tokenizer_with_bos
 
@@ -38,6 +37,11 @@ def boot(
     Returns:
         The bridge to the loaded model.
     """
+    # Lazy import to avoid circular import
+    from transformer_lens.factories.architecture_adapter_factory import (
+        ArchitectureAdapterFactory,
+    )
+
     hf_config = AutoConfig.from_pretrained(model_name, **kwargs)
     adapter = ArchitectureAdapterFactory.select_architecture_adapter(hf_config)
     default_config = adapter.default_cfg
@@ -125,3 +129,6 @@ def setup_tokenizer(
         tokenizer.bos_token = tokenizer.eos_token
 
     return tokenizer
+
+
+setattr(TransformerBridge, "boot_transformers", staticmethod(boot))