Feat/external architecture registration by huseyincavusbi · Pull Request #1307 · TransformerLensOrg/TransformerLens

huseyincavusbi · 2026-05-16T11:05:36Z

Hi @jlarson4,

Description

Three changes to allow users to register custom architecture adapters without forking TransformerLens:

register_adapter() method — Classmethod on ArchitectureAdapterFactory that adds custom architectures to the registry at runtime
Entry-point discovery — discover_entry_points() scans installed packages for adapters declared via transformer_lens.architectures entry group, auto-called on first select_architecture_adapter()
Documentation — New page under Contributing with examples for runtime registration, entry-point declaration, and example package layout

Closes #1298

Type of change

New feature (non-breaking change which adds functionality)
This change requires a documentation update

Checklist

I have commented my code, particularly in hard-to-understand areas
My changes generate no new warnings
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes
I have not rewritten tests relating to key interfaces which would affect backward compatibility

…overy

…oint tests

Copilot

Pull request overview

Adds a public mechanism for registering custom architecture adapters in TransformerLens without forking the project, plus accompanying tests and documentation.

Changes:

New ArchitectureAdapterFactory.register_adapter() classmethod for runtime registration, and discover_entry_points() that reads adapters from the transformer_lens.architectures entry-point group on first use.
New unit test module covering registration, discovery idempotency, and select error paths.
New documentation page under Contributing describing both registration paths and an example external package layout.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
transformer_lens/factories/architecture_adapter_factory.py	Adds `register_adapter` and `discover_entry_points` classmethods; calls discovery at the start of `select_architecture_adapter`.
tests/unit/model_bridge/test_architecture_adapter_factory.py	New tests for `SUPPORTED_ARCHITECTURES` contents, runtime registration, error handling, and discovery idempotency.
docs/source/content/contributing.md	Adds the new doc page to the toctree.
docs/source/content/adapter_development/external-adapter-registration.md	New doc page explaining runtime and entry-point registration, with example package layout.

Comments suppressed due to low confidence (1)

tests/unit/model_bridge/test_architecture_adapter_factory.py:58

This test does not actually verify "overwrite" behavior. It registers MockArchitectureAdapter twice for the same key, so the post-condition cls._adapters[key] is first is trivially true regardless of whether the second register_adapter call overwrote anything. To meaningfully test overwrite behavior, register a second, different adapter class and assert that the registry now points at the new class (and not at first).

    def test_register_overwrites_existing(self):
        key = "TestOverwriteForCausalLM"
        ArchitectureAdapterFactory.register_adapter(key, MockArchitectureAdapter)
        first = ArchitectureAdapterFactory._adapters[key]
        ArchitectureAdapterFactory.register_adapter(key, MockArchitectureAdapter)
        assert ArchitectureAdapterFactory._adapters[key] is first

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+class TestRegisterAdapter:
+    """Verify runtime adapter registration."""
+
+    def test_register_adds_to_adapters(self):
+        key = "TestMockForCausalLM"
+        ArchitectureAdapterFactory.register_adapter(key, MockArchitectureAdapter)
+        assert key in ArchitectureAdapterFactory._adapters
+
+    def test_register_overwrites_existing(self):
+        key = "TestOverwriteForCausalLM"
+        ArchitectureAdapterFactory.register_adapter(key, MockArchitectureAdapter)
+        first = ArchitectureAdapterFactory._adapters[key]
+        ArchitectureAdapterFactory.register_adapter(key, MockArchitectureAdapter)
+        assert ArchitectureAdapterFactory._adapters[key] is first
+
+    def test_select_returns_registered_adapter(self):
+        key = "TestSelectForCausalLM"
+        ArchitectureAdapterFactory.register_adapter(key, MockArchitectureAdapter)
+        cfg = _make_cfg(architecture=key)
+        adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+        assert isinstance(adapter, MockArchitectureAdapter)
+
+
+class TestSelectErrors:
+    """Verify error handling in select_architecture_adapter."""
+
+    def test_unknown_architecture_raises(self):
+        cfg = _make_cfg(architecture="NonExistentForCausalLM")
+        with pytest.raises(ValueError, match="Unsupported architecture"):
+            ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+
+    def test_none_architecture_raises(self):
+        cfg = _make_cfg(architecture=None)
+        with pytest.raises(ValueError, match="must have architecture set"):
+            ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+
+
+class TestDiscoverEntryPoints:
+    """Verify entry-point discovery behavior."""
+
+    def test_discover_is_idempotent(self):
+        ArchitectureAdapterFactory._entry_points_discovered = False
+        ArchitectureAdapterFactory.discover_entry_points()
+        first_run = ArchitectureAdapterFactory._entry_points_discovered
+        ArchitectureAdapterFactory.discover_entry_points()
+        assert ArchitectureAdapterFactory._entry_points_discovered is first_run is True
+
+    def test_discover_does_not_remove_existing(self):
+        key = "TestPreserveForCausalLM"
+        ArchitectureAdapterFactory.register_adapter(key, MockArchitectureAdapter)
+        ArchitectureAdapterFactory._entry_points_discovered = False
+        ArchitectureAdapterFactory.discover_entry_points()
+        assert key in ArchitectureAdapterFactory._adapters


jlarson4 · 2026-05-18T14:37:29Z

+    discovery of adapters from installed packages via entry points.
+    """

    _adapters = SUPPORTED_ARCHITECTURES


This is a great change, thanks for implementing @huseyincavusbi

… class

jlarson4 · 2026-05-18T14:24:30Z

+        else:
+            for ep in eps:
+                try:
+                    cls._adapters[ep.name] = ep.load()


We should add a guard here:

If an adapter already exists as part of the main TransformerLens package, we do not want it to be automatically overwritten by a custom adapter from another package. Let's check for the existence of ep.name in our _adapters list, and log a warning.

For the warning, I am thinking something along the lines of "Custom architecture adapter {ep.name} provided by {ep.dist.name} attempted to override a native adapter. If you'd like to use this custom adapter, register it explicitly with register_adapter".

jlarson4 · 2026-05-18T14:25:28Z

+            ...     def __init__(self, cfg):
+            ...         super().__init__(cfg)
+            >>> ArchitectureAdapterFactory.register_adapter("MyModelForCausalLM", MyAdapter)
+            >>> "MyModelForCausalLM" in ArchitectureAdapterFactory._adapters


This example references the private cls._adapters attribute. Better to demonstrate the public path via select_architecture_adapter.

jlarson4 · 2026-05-18T14:28:04Z

+    return TransformerBridgeConfig(**defaults)
+
+
+class TestSupportedArchitectures:


Drop this test

jlarson4 · 2026-05-18T14:32:46Z

There are no tests included for testing the entry point discovery system. Would it be possible to add some additional tests covering those features? Specifically we want to cover:

Successfully entry point discovery

Override prevention

Failed to register bad adapter

jlarson4 · 2026-05-18T14:36:54Z

+
+> **Important:** The architecture name you register (e.g. `"MyModelForCausalLM"`) must match the `architectures` field in the model's HuggingFace `config.json`. TransformerLens reads this field to look up the adapter.
+
+### 2. Entry-point registration (recommended for packages)


Add a Note at the end of this section documenting the warning behavior that occurs if you try to override a native adapter via entry-point

jlarson4 · 2026-05-18T14:38:49Z

@huseyincavusbi just a couple small comments related to entry point behaviors, all together this is a great implementation, thank you for taking this on. Feel free to ping me if you have any questions

jlarson4 · 2026-05-20T00:01:14Z

Awesome work on this @huseyincavusbi! I am approving and merging now, this will be included in the next release

huseyincavusbi added 3 commits May 12, 2026 08:56

feat: External architecture adapter registration and entry-point disc…

6af217c

…overy

docs: External adapter registration guide with examples

a68a198

test: ArchitectureAdapterFactory registration, selection, and entry-p…

ba01882

…oint tests

Copilot AI review requested due to automatic review settings May 16, 2026 11:05

Copilot started reviewing on behalf of huseyincavusbi May 16, 2026 11:06 View session

Copilot AI reviewed May 16, 2026

View reviewed changes

huseyincavusbi added 4 commits May 16, 2026 13:08

fix: Clarify matching requirement and add warnings for failed discovery

93c37f7

fix: Make register_adapter doctest self-contained with inline adapter…

2ef3161

… class

fix: Address Copilot review - alias, loop isolation, test state leak

856c1f2

fix: Apply black formatting to factory file

12fd1fa

jlarson4 reviewed May 18, 2026

View reviewed changes

huseyincavusbi added 8 commits May 19, 2026 22:49

fix: use select_architecture_adapter in register_adapter doctest

f728287

fix: guard native adapter override in discover_entry_points

c1a89d8

test: drop TestSupportedArchitectures

b676971

test: add entry-point discovery tests

aa6169c

docs: document native adapter override prevention

f726c08

fix: apply black formatting to test file

49e365e

fix: handle None ep.dist in entry-point guard warning

6ac1e97

fix: apply black formatting to factory file

a2ae191

jlarson4 merged commit a152ccc into TransformerLensOrg:dev May 20, 2026
24 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feat/external architecture registration#1307

Feat/external architecture registration#1307
jlarson4 merged 15 commits into
TransformerLensOrg:devfrom
huseyincavusbi:feat/external-architecture-registration

huseyincavusbi commented May 16, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

jlarson4 May 18, 2026

Uh oh!

jlarson4 May 18, 2026

Uh oh!

jlarson4 May 18, 2026

Uh oh!

jlarson4 May 18, 2026

Uh oh!

jlarson4 May 18, 2026

Uh oh!

jlarson4 May 18, 2026

Uh oh!

jlarson4 commented May 18, 2026

Uh oh!

jlarson4 commented May 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return TransformerBridgeConfig(**defaults)


		class TestSupportedArchitectures:


		> Important: The architecture name you register (e.g. `"MyModelForCausalLM"`) must match the `architectures` field in the model's HuggingFace `config.json`. TransformerLens reads this field to look up the adapter.

		### 2. Entry-point registration (recommended for packages)

Conversation

huseyincavusbi commented May 16, 2026

Description

Type of change

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

jlarson4 May 18, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 18, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 18, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 18, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 18, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 18, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 commented May 18, 2026

Uh oh!

jlarson4 commented May 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants