created new base config class (#1042)

* created new base config class

* cleaned up imports

* reorganized config

* setup transformer bridge config

* fixed docstring issue

* fixed typing issues

* ran format

* fixed docstring

* fixed import issues

* ran format

* fixed type checking again

* fixed import again

* fixed doc string

* removed import

* seperated devices

* cleaned up utils a bit

* ran format

* fixed import

* fixed imports

* ran format

* fixed typing

* updated typing

* updated name

* changed to python 3.12

* cleaned up comments

* removed extra functions

Author: bryce13950

.github/workflows/checks.yml

@@ -99,7 +99,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.13"
+          python-version: "3.12"
           cache: "poetry"
       - name: Cache Models used with Tests
         uses: actions/cache@v3

tests/acceptance/test_multi_gpu.py

@@ -4,7 +4,7 @@
 import torch
 
 from transformer_lens.HookedTransformer import HookedTransformer
-from transformer_lens.utilities.devices import get_best_available_device
+from transformer_lens.utilities.multi_gpu import get_best_available_device
 
 
 @pytest.fixture

tests/integration/test_attention_mask.py

@@ -1,8 +1,8 @@
 import torch
 
 from transformer_lens import utils
+from transformer_lens.config import HookedTransformerConfig
 from transformer_lens.HookedTransformer import HookedTransformer
-from transformer_lens.HookedTransformerConfig import HookedTransformerConfig
 
 
 def test_attention_mask():

tests/integration/test_grouped_query_attention.py

@@ -3,7 +3,7 @@
 
 from transformer_lens import HookedTransformer
 from transformer_lens.components import Attention, GroupedQueryAttention
-from transformer_lens.HookedTransformerConfig import HookedTransformerConfig
+from transformer_lens.config import HookedTransformerConfig
 
 
 def test_grouped_query_attention_output_is_correct():

tests/mocks/architecture_adapter.py

@@ -4,6 +4,7 @@
 import pytest
 import torch.nn as nn
 
+from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.generalized_components import (
     AttentionBridge,
@@ -19,8 +20,17 @@ class MockArchitectureAdapter(ArchitectureAdapter):
 
     def __init__(self, cfg=None):
         if cfg is None:
-            # Create a minimal config for testing
-            cfg = SimpleNamespace(d_mlp=512, intermediate_size=512, default_prepend_bos=True)
+            # Create a minimal TransformerBridgeConfig for testing
+            cfg = TransformerBridgeConfig(
+                d_model=512,
+                d_head=64,
+                n_layers=2,
+                n_ctx=1024,
+                d_vocab=1000,
+                d_mlp=2048,
+                default_prepend_bos=True,
+                architecture="GPT2LMHeadModel",  # Default test architecture
+            )
         super().__init__(cfg)
         # Use actual bridge instances instead of tuples
         # Provide minimal config to components that require it

tests/unit/components/test_attention.py

@@ -5,7 +5,7 @@
 from transformers.utils import is_bitsandbytes_available
 
 from transformer_lens.components import Attention
-from transformer_lens.HookedTransformerConfig import HookedTransformerConfig
+from transformer_lens.config import HookedTransformerConfig
 from transformer_lens.utilities.attention import complex_attn_linear
 
 if is_bitsandbytes_available():

tests/unit/factories/test_activation_function_factory.py

@@ -1,10 +1,10 @@
 import pytest
 import torch
 
+from transformer_lens.config import HookedTransformerConfig
 from transformer_lens.factories.activation_function_factory import (
     ActivationFunctionFactory,
 )
-from transformer_lens.HookedTransformerConfig import HookedTransformerConfig
 from transformer_lens.utilities.activation_functions import SUPPORTED_ACTIVATIONS
 
 

tests/unit/factories/test_mlp_factory.py

@@ -4,8 +4,8 @@
 from transformer_lens.components.mlps.gated_mlp import GatedMLP
 from transformer_lens.components.mlps.gated_mlp_4bit import GatedMLP4Bit
 from transformer_lens.components.mlps.mlp import MLP
+from transformer_lens.config import HookedTransformerConfig
 from transformer_lens.factories.mlp_factory import MLPFactory
-from transformer_lens.HookedTransformerConfig import HookedTransformerConfig
 
 
 def test_create_mlp_basic():

tests/unit/pretrained_weight_conversions/test_neo.py

@@ -3,7 +3,7 @@
 import torch
 
 from transformer_lens import HookedTransformer
-from transformer_lens.HookedTransformerConfig import HookedTransformerConfig
+from transformer_lens.config import HookedTransformerConfig
 from transformer_lens.pretrained.weight_conversions.neo import convert_neo_weights
 
 

tests/unit/test_architecture_adapter.py

@@ -9,6 +9,7 @@
     mock_model_adapter,
 )
 from tests.mocks.models import MockGemma3Model
+from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge.supported_architectures.gemma3 import (
     Gemma3ArchitectureAdapter,
 )
@@ -37,17 +38,20 @@ def test_get_remote_component_with_mock(
     assert isinstance(mlp, nn.Module)
 
 
-class DummyHFConfig:
-    def __init__(self):
-        self.num_attention_heads = 8
-        self.num_key_value_heads = 8
-        self.hidden_size = 128
-        # Add any other attributes needed by the adapter here
-
-
 @pytest.fixture
 def cfg():
-    return DummyHFConfig()
+    return TransformerBridgeConfig(
+        d_model=128,
+        d_head=16,  # 128 / 8 heads
+        n_layers=2,
+        n_ctx=1024,
+        n_heads=8,
+        d_vocab=1000,
+        d_mlp=512,
+        n_key_value_heads=8,
+        default_prepend_bos=True,
+        architecture="Gemma3ForCausalLM",  # Test architecture
+    )
 
 
 @pytest.fixture