feat: Bridge.boot should allow using alias model names, but show a deprecation warning (#1028)

* Automatically replace aliased model name and show deprecation warning

* add test for aliased model name and deprecation

Author: hijohnnylin

tests/integration/model_bridge/test_bridge_integration.py

@@ -4,6 +4,8 @@
 including model initialization, text generation, hooks, and caching.
 """
 
+import logging
+
 import pytest
 import torch
 
@@ -21,6 +23,35 @@ def test_model_initialization():
     assert isinstance(bridge.original_model, torch.nn.Module), "Model should be a PyTorch module"
 
 
+def test_model_initialization_with_alias(caplog):
+    """Test that the model can be initialized correctly with an alias and logs deprecation warning."""
+
+    model_name = "gpt2-small"
+
+    # Set logging level to capture warnings
+    with caplog.at_level(logging.WARNING):
+        bridge = TransformerBridge.boot_transformers(model_name)
+
+        # Basic assertions
+        assert bridge is not None, "Bridge should be initialized"
+        assert bridge.tokenizer is not None, "Tokenizer should be initialized"
+        assert isinstance(
+            bridge.original_model, torch.nn.Module
+        ), "Model should be a PyTorch module"
+
+    # Check that a deprecation warning was logged
+    deprecation_found = False
+    for record in caplog.records:
+        if "DEPRECATED" in record.message:
+            deprecation_found = True
+            # Verify the warning contains expected content
+            assert "gpt2-small" in record.message, "Warning should mention the deprecated alias"
+            assert "gpt2" in record.message, "Warning should mention the official name"
+            break
+
+    assert deprecation_found, "Expected deprecation warning for alias 'gpt2-small' was not logged"
+
+
 def test_text_generation():
     """Test basic text generation functionality."""
     model_name = "gpt2"  # Use a smaller model for testing

transformer_lens/model_bridge/sources/transformers.py

@@ -3,8 +3,8 @@
 This module provides functionality to load and convert models from HuggingFace to TransformerLens format.
 """
 
-
 import copy
+import logging
 import os
 
 import torch
@@ -16,6 +16,7 @@
 )
 
 from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.supported_models import MODEL_ALIASES
 from transformer_lens.utils import get_tokenizer_with_bos
 
 
@@ -111,6 +112,16 @@ def boot(
         ArchitectureAdapterFactory,
     )
 
+    # MODEL_ALIASES is a dict of {official_name: [alias1, alias2, ...]}
+    # Check if model_name that the user passed is an alias, and if so, use the official name
+    for official_name, aliases in MODEL_ALIASES.items():
+        if model_name in aliases:
+            logging.warning(
+                f"DEPRECATED: You are using a deprecated, model_name alias '{model_name}'. TransformerLens will now load the official transformers model name, '{official_name}' instead.\n Please update your code to use the official name by changing model_name from '{model_name}' to '{official_name}'.\nSince TransformerLens v3, all model names should be the official transformers model names.\nThe aliases will be removed in the next version of TransformerLens, so please do the update now."
+            )
+            model_name = official_name
+            break
+
     hf_config = AutoConfig.from_pretrained(model_name, output_attentions=True)
 
     # Apply config variables to hf_config before selecting adapter