Coverage for transformer_lens/factories/architecture_adapter_factory.py: 96%
38 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Architecture adapter factory.
3This module provides a factory for creating architecture adapters, including
4support for external registration and entry-point discovery.
5"""
7import warnings
8from importlib.metadata import entry_points
10from transformer_lens.config import TransformerBridgeConfig
11from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
12from transformer_lens.model_bridge.supported_architectures import (
13 AfmoeArchitectureAdapter,
14 ApertusArchitectureAdapter,
15 ArceeArchitectureAdapter,
16 ASTArchitectureAdapter,
17 AudioFlamingo3ArchitectureAdapter,
18 BaichuanArchitectureAdapter,
19 BambaArchitectureAdapter,
20 BartArchitectureAdapter,
21 BD3LMArchitectureAdapter,
22 BertArchitectureAdapter,
23 BitNetArchitectureAdapter,
24 BlenderbotArchitectureAdapter,
25 BloomArchitectureAdapter,
26 CodeGenArchitectureAdapter,
27 Cohere2ArchitectureAdapter,
28 CohereArchitectureAdapter,
29 DeepSeekV2ArchitectureAdapter,
30 DeepSeekV3ArchitectureAdapter,
31 DeepSeekV4ArchitectureAdapter,
32 DreamArchitectureAdapter,
33 Emu3ArchitectureAdapter,
34 Ernie4_5_MoeArchitectureAdapter,
35 Ernie4_5ArchitectureAdapter,
36 Exaone4ArchitectureAdapter,
37 ExaoneArchitectureAdapter,
38 FalconArchitectureAdapter,
39 FalconH1ArchitectureAdapter,
40 FalconMambaArchitectureAdapter,
41 FlexOlmoArchitectureAdapter,
42 Florence2ArchitectureAdapter,
43 Gemma1ArchitectureAdapter,
44 Gemma2ArchitectureAdapter,
45 Gemma3ArchitectureAdapter,
46 Gemma3MultimodalArchitectureAdapter,
47 Gemma3nArchitectureAdapter,
48 Gemma4ArchitectureAdapter,
49 Gemma4TextArchitectureAdapter,
50 GiddArchitectureAdapter,
51 Glm4ArchitectureAdapter,
52 Glm4MoeArchitectureAdapter,
53 Glm4MoeLiteArchitectureAdapter,
54 Glm4vArchitectureAdapter,
55 GlmArchitectureAdapter,
56 GlmAsrArchitectureAdapter,
57 GlmMoeDsaArchitectureAdapter,
58 GPT2ArchitectureAdapter,
59 Gpt2LmHeadCustomArchitectureAdapter,
60 GPTBigCodeArchitectureAdapter,
61 GptjArchitectureAdapter,
62 GPTOSSArchitectureAdapter,
63 GraniteArchitectureAdapter,
64 GraniteMoeArchitectureAdapter,
65 GraniteMoeHybridArchitectureAdapter,
66 HrmTextArchitectureAdapter,
67 HubertArchitectureAdapter,
68 HunYuanDenseV1ArchitectureAdapter,
69 HyenaDNAArchitectureAdapter,
70 Idefics3ArchitectureAdapter,
71 InternLM2ArchitectureAdapter,
72 Jais2ArchitectureAdapter,
73 JambaArchitectureAdapter,
74 JetMoeArchitectureAdapter,
75 LagunaArchitectureAdapter,
76 LEDArchitectureAdapter,
77 Lfm2ArchitectureAdapter,
78 Lfm2MoeArchitectureAdapter,
79 LLaDA2MoeArchitectureAdapter,
80 LLaDAArchitectureAdapter,
81 Llama4ArchitectureAdapter,
82 Llama4MultimodalArchitectureAdapter,
83 LlamaArchitectureAdapter,
84 LlavaArchitectureAdapter,
85 LlavaNextArchitectureAdapter,
86 LlavaOnevisionArchitectureAdapter,
87 LongT5ArchitectureAdapter,
88 M2M100ArchitectureAdapter,
89 Mamba2ArchitectureAdapter,
90 MambaArchitectureAdapter,
91 MarianArchitectureAdapter,
92 MBartArchitectureAdapter,
93 MingptArchitectureAdapter,
94 MiniMaxM2ArchitectureAdapter,
95 Ministral3ArchitectureAdapter,
96 Mistral3ArchitectureAdapter,
97 MistralArchitectureAdapter,
98 MixtralArchitectureAdapter,
99 ModernBertDecoderArchitectureAdapter,
100 MPTArchitectureAdapter,
101 MuseGlimmerArchitectureAdapter,
102 MusicFlamingoArchitectureAdapter,
103 NanoChatArchitectureAdapter,
104 NanogptArchitectureAdapter,
105 NativeArchitectureAdapter,
106 NeelSoluOldArchitectureAdapter,
107 NemotronArchitectureAdapter,
108 NemotronHArchitectureAdapter,
109 NeoArchitectureAdapter,
110 NeoxArchitectureAdapter,
111 Olmo2ArchitectureAdapter,
112 Olmo3ArchitectureAdapter,
113 OlmoArchitectureAdapter,
114 OlmoeArchitectureAdapter,
115 OlmoHybridArchitectureAdapter,
116 OpenAIGPTArchitectureAdapter,
117 OpenElmArchitectureAdapter,
118 OptArchitectureAdapter,
119 OuroArchitectureAdapter,
120 PegasusArchitectureAdapter,
121 Phi3ArchitectureAdapter,
122 PhiArchitectureAdapter,
123 PhiMoEArchitectureAdapter,
124 PretrainArchitectureAdapter,
125 Qwen2_5_VLArchitectureAdapter,
126 Qwen2ArchitectureAdapter,
127 Qwen2AudioArchitectureAdapter,
128 Qwen2MoeArchitectureAdapter,
129 Qwen3_5ArchitectureAdapter,
130 Qwen3_5MoeArchitectureAdapter,
131 Qwen3_5MoeMultimodalArchitectureAdapter,
132 Qwen3_5MultimodalArchitectureAdapter,
133 Qwen3ArchitectureAdapter,
134 Qwen3MoeArchitectureAdapter,
135 Qwen3NextArchitectureAdapter,
136 Qwen3VLArchitectureAdapter,
137 Qwen3VLMoeArchitectureAdapter,
138 QwenArchitectureAdapter,
139 RavenArchitectureAdapter,
140 RecurrentGemmaArchitectureAdapter,
141 RobertaArchitectureAdapter,
142 RWKV7ArchitectureAdapter,
143 RwkvArchitectureAdapter,
144 SeedOssArchitectureAdapter,
145 SmolLM3ArchitectureAdapter,
146 StableLmArchitectureAdapter,
147 Starcoder2ArchitectureAdapter,
148 SwitchTransformersArchitectureAdapter,
149 T5ArchitectureAdapter,
150 T5Gemma2ArchitectureAdapter,
151 T5GemmaArchitectureAdapter,
152 VaultGemmaArchitectureAdapter,
153 ViTArchitectureAdapter,
154 Wav2Vec2ArchitectureAdapter,
155 XGLMArchitectureAdapter,
156 YoutuArchitectureAdapter,
157 Zamba2ArchitectureAdapter,
158)
160# Export supported architectures
161SUPPORTED_ARCHITECTURES = {
162 "AfmoeForCausalLM": AfmoeArchitectureAdapter,
163 "ApertusForCausalLM": ApertusArchitectureAdapter,
164 "ArceeForCausalLM": ArceeArchitectureAdapter,
165 "ASTForAudioClassification": ASTArchitectureAdapter,
166 "BaiChuanForCausalLM": BaichuanArchitectureAdapter,
167 "BaichuanForCausalLM": BaichuanArchitectureAdapter,
168 "BambaForCausalLM": BambaArchitectureAdapter,
169 "BartForConditionalGeneration": BartArchitectureAdapter,
170 "BD3LM": BD3LMArchitectureAdapter,
171 "DreamModel": DreamArchitectureAdapter,
172 "Emu3ForConditionalGeneration": Emu3ArchitectureAdapter,
173 "AudioFlamingo3ForConditionalGeneration": AudioFlamingo3ArchitectureAdapter,
174 "FlexOlmoForCausalLM": FlexOlmoArchitectureAdapter,
175 "GiddForDiffusionLM": GiddArchitectureAdapter,
176 "HyenaDNAForCausalLM": HyenaDNAArchitectureAdapter,
177 "LLaDA2MoeModelLM": LLaDA2MoeArchitectureAdapter,
178 "Jais2ForCausalLM": Jais2ArchitectureAdapter,
179 "JetMoeForCausalLM": JetMoeArchitectureAdapter,
180 # jetmoe-8b checkpoints predate the native port and use the remote-code capitalization
181 "JetMoEForCausalLM": JetMoeArchitectureAdapter,
182 "LagunaForCausalLM": LagunaArchitectureAdapter,
183 "Ministral3ForCausalLM": Ministral3ArchitectureAdapter,
184 "VaultGemmaForCausalLM": VaultGemmaArchitectureAdapter,
185 "YoutuForCausalLM": YoutuArchitectureAdapter,
186 "ModernBertDecoderForCausalLM": ModernBertDecoderArchitectureAdapter,
187 "MusicFlamingoForConditionalGeneration": MusicFlamingoArchitectureAdapter,
188 "NanoChatForCausalLM": NanoChatArchitectureAdapter,
189 "RwkvForCausalLM": RwkvArchitectureAdapter,
190 "SwitchTransformersForConditionalGeneration": SwitchTransformersArchitectureAdapter,
191 "BertForMaskedLM": BertArchitectureAdapter,
192 "BertLMHeadModel": BertArchitectureAdapter,
193 "RobertaForMaskedLM": RobertaArchitectureAdapter,
194 "BitNetForCausalLM": BitNetArchitectureAdapter,
195 "BlenderbotForConditionalGeneration": BlenderbotArchitectureAdapter,
196 "BloomForCausalLM": BloomArchitectureAdapter,
197 "BloomModel": BloomArchitectureAdapter,
198 "CodeGenForCausalLM": CodeGenArchitectureAdapter,
199 "Cohere2ForCausalLM": Cohere2ArchitectureAdapter,
200 "CohereForCausalLM": CohereArchitectureAdapter,
201 "DeepseekV2ForCausalLM": DeepSeekV2ArchitectureAdapter,
202 "DeepseekV3ForCausalLM": DeepSeekV3ArchitectureAdapter,
203 "Ernie4_5ForCausalLM": Ernie4_5ArchitectureAdapter,
204 "Ernie4_5_MoeForCausalLM": Ernie4_5_MoeArchitectureAdapter,
205 "ExaoneForCausalLM": ExaoneArchitectureAdapter,
206 "Exaone4ForCausalLM": Exaone4ArchitectureAdapter,
207 "DeepseekV4ForCausalLM": DeepSeekV4ArchitectureAdapter,
208 "FalconForCausalLM": FalconArchitectureAdapter,
209 "FalconH1ForCausalLM": FalconH1ArchitectureAdapter,
210 "FalconMambaForCausalLM": FalconMambaArchitectureAdapter,
211 "Florence2ForConditionalGeneration": Florence2ArchitectureAdapter,
212 "GemmaForCausalLM": Gemma1ArchitectureAdapter, # Default to Gemma1 as it's the original version
213 "Gemma1ForCausalLM": Gemma1ArchitectureAdapter,
214 "Gemma2ForCausalLM": Gemma2ArchitectureAdapter,
215 "Gemma3ForCausalLM": Gemma3ArchitectureAdapter,
216 "Gemma3ForConditionalGeneration": Gemma3MultimodalArchitectureAdapter,
217 "Gemma3nForConditionalGeneration": Gemma3nArchitectureAdapter,
218 "Gemma4ForConditionalGeneration": Gemma4ArchitectureAdapter,
219 # The unified (encoder-free) 12B variant's text decoder is a strict structural
220 # subset of gemma4 (no PLE, no MoE — both optional in the adapter), with the
221 # same module paths. Requires transformers >= 5.10 to load.
222 "Gemma4UnifiedForConditionalGeneration": Gemma4ArchitectureAdapter,
223 "Gemma4ForCausalLM": Gemma4TextArchitectureAdapter,
224 "GraniteForCausalLM": GraniteArchitectureAdapter,
225 "GraniteMoeForCausalLM": GraniteMoeArchitectureAdapter,
226 "GraniteMoeHybridForCausalLM": GraniteMoeHybridArchitectureAdapter,
227 "GlmForCausalLM": GlmArchitectureAdapter,
228 "Glm4ForCausalLM": Glm4ArchitectureAdapter,
229 "Glm4vForConditionalGeneration": Glm4vArchitectureAdapter,
230 "GlmMoeDsaForCausalLM": GlmMoeDsaArchitectureAdapter,
231 "Glm4MoeForCausalLM": Glm4MoeArchitectureAdapter,
232 "Glm4MoeLiteForCausalLM": Glm4MoeLiteArchitectureAdapter,
233 "GlmAsrForConditionalGeneration": GlmAsrArchitectureAdapter,
234 "GPT2LMHeadModel": GPT2ArchitectureAdapter,
235 "GPTBigCodeForCausalLM": GPTBigCodeArchitectureAdapter,
236 "GptOssForCausalLM": GPTOSSArchitectureAdapter,
237 "GPT2LMHeadCustomModel": Gpt2LmHeadCustomArchitectureAdapter,
238 "GPTJForCausalLM": GptjArchitectureAdapter,
239 "HrmTextForCausalLM": HrmTextArchitectureAdapter,
240 "HubertForCTC": HubertArchitectureAdapter,
241 "HubertModel": HubertArchitectureAdapter,
242 "Wav2Vec2ForCTC": Wav2Vec2ArchitectureAdapter,
243 "Wav2Vec2ForPreTraining": Wav2Vec2ArchitectureAdapter,
244 "Wav2Vec2Model": Wav2Vec2ArchitectureAdapter,
245 "HunYuanDenseV1ForCausalLM": HunYuanDenseV1ArchitectureAdapter,
246 "Idefics3ForConditionalGeneration": Idefics3ArchitectureAdapter,
247 "InternLM2ForCausalLM": InternLM2ArchitectureAdapter,
248 "JambaForCausalLM": JambaArchitectureAdapter,
249 "LEDForConditionalGeneration": LEDArchitectureAdapter,
250 "LLaDAModelLM": LLaDAArchitectureAdapter,
251 "LlamaForCausalLM": LlamaArchitectureAdapter,
252 "Llama4ForCausalLM": Llama4ArchitectureAdapter,
253 "Llama4ForConditionalGeneration": Llama4MultimodalArchitectureAdapter,
254 "LlavaForConditionalGeneration": LlavaArchitectureAdapter,
255 "LlavaNextForConditionalGeneration": LlavaNextArchitectureAdapter,
256 "LlavaOnevisionForConditionalGeneration": LlavaOnevisionArchitectureAdapter,
257 "Lfm2ForCausalLM": Lfm2ArchitectureAdapter,
258 "Lfm2MoeForCausalLM": Lfm2MoeArchitectureAdapter,
259 "LongT5ForConditionalGeneration": LongT5ArchitectureAdapter,
260 "M2M100ForConditionalGeneration": M2M100ArchitectureAdapter,
261 "Mamba2ForCausalLM": Mamba2ArchitectureAdapter,
262 "MambaForCausalLM": MambaArchitectureAdapter,
263 "MarianMTModel": MarianArchitectureAdapter,
264 "MBartForConditionalGeneration": MBartArchitectureAdapter,
265 "MiniMaxM2ForCausalLM": MiniMaxM2ArchitectureAdapter,
266 "NemotronForCausalLM": NemotronArchitectureAdapter,
267 "NemotronHForCausalLM": NemotronHArchitectureAdapter,
268 "MixtralForCausalLM": MixtralArchitectureAdapter,
269 "MistralForCausalLM": MistralArchitectureAdapter,
270 "Mistral3ForConditionalGeneration": Mistral3ArchitectureAdapter,
271 "MPTForCausalLM": MPTArchitectureAdapter,
272 "MptForCausalLM": MPTArchitectureAdapter,
273 "MuseGlimmerForConditionalGeneration": MuseGlimmerArchitectureAdapter,
274 "NeoForCausalLM": NeoArchitectureAdapter,
275 "NeoXForCausalLM": NeoxArchitectureAdapter,
276 "NeelSoluOldForCausalLM": NeelSoluOldArchitectureAdapter,
277 "OlmoForCausalLM": OlmoArchitectureAdapter,
278 "Olmo2ForCausalLM": Olmo2ArchitectureAdapter,
279 "Olmo3ForCausalLM": Olmo3ArchitectureAdapter,
280 "OlmoeForCausalLM": OlmoeArchitectureAdapter,
281 "OlmoHybridForCausalLM": OlmoHybridArchitectureAdapter,
282 "OpenAIGPTLMHeadModel": OpenAIGPTArchitectureAdapter,
283 "OpenELMForCausalLM": OpenElmArchitectureAdapter,
284 "OPTForCausalLM": OptArchitectureAdapter,
285 "PegasusForConditionalGeneration": PegasusArchitectureAdapter,
286 "OuroForCausalLM": OuroArchitectureAdapter,
287 "PhiForCausalLM": PhiArchitectureAdapter,
288 "Phi3ForCausalLM": Phi3ArchitectureAdapter,
289 "PhiMoEForCausalLM": PhiMoEArchitectureAdapter,
290 "QwenForCausalLM": QwenArchitectureAdapter,
291 "Qwen2ForCausalLM": Qwen2ArchitectureAdapter,
292 "Qwen2_5_VLForConditionalGeneration": Qwen2_5_VLArchitectureAdapter,
293 "Qwen2AudioForConditionalGeneration": Qwen2AudioArchitectureAdapter,
294 "Qwen2MoeForCausalLM": Qwen2MoeArchitectureAdapter,
295 "Qwen3ForCausalLM": Qwen3ArchitectureAdapter,
296 "Qwen3VLForConditionalGeneration": Qwen3VLArchitectureAdapter,
297 "Qwen3VLMoeForConditionalGeneration": Qwen3VLMoeArchitectureAdapter,
298 "Qwen3MoeForCausalLM": Qwen3MoeArchitectureAdapter,
299 "Qwen3NextForCausalLM": Qwen3NextArchitectureAdapter,
300 "Qwen3_5ForCausalLM": Qwen3_5ArchitectureAdapter,
301 "Qwen3_5ForConditionalGeneration": Qwen3_5MultimodalArchitectureAdapter,
302 "Qwen3_5MoeForCausalLM": Qwen3_5MoeArchitectureAdapter,
303 "Qwen3_5MoeForConditionalGeneration": Qwen3_5MoeMultimodalArchitectureAdapter,
304 "RavenForCausalLM": RavenArchitectureAdapter,
305 "RecurrentGemmaForCausalLM": RecurrentGemmaArchitectureAdapter,
306 "SeedOssForCausalLM": SeedOssArchitectureAdapter,
307 "RWKV7ForCausalLM": RWKV7ArchitectureAdapter,
308 "SmolLM3ForCausalLM": SmolLM3ArchitectureAdapter,
309 "StableLmForCausalLM": StableLmArchitectureAdapter,
310 "Starcoder2ForCausalLM": Starcoder2ArchitectureAdapter,
311 "T5ForConditionalGeneration": T5ArchitectureAdapter,
312 "MT5ForConditionalGeneration": T5ArchitectureAdapter,
313 "T5WithLMHeadModel": T5ArchitectureAdapter,
314 "T5GemmaForConditionalGeneration": T5GemmaArchitectureAdapter,
315 "T5Gemma2ForConditionalGeneration": T5Gemma2ArchitectureAdapter,
316 "XGLMForCausalLM": XGLMArchitectureAdapter,
317 "Zamba2ForCausalLM": Zamba2ArchitectureAdapter,
318 "NanoGPTForCausalLM": NanogptArchitectureAdapter,
319 "TransformerLensNative": NativeArchitectureAdapter,
320 "TransformerLensPretrain": PretrainArchitectureAdapter,
321 "MinGPTForCausalLM": MingptArchitectureAdapter,
322 "GPTNeoForCausalLM": NeoArchitectureAdapter,
323 "GPTNeoXForCausalLM": NeoxArchitectureAdapter,
324 "ViTModel": ViTArchitectureAdapter,
325 "ViTForImageClassification": ViTArchitectureAdapter,
326 "DeiTModel": ViTArchitectureAdapter,
327 "DeiTForImageClassification": ViTArchitectureAdapter, # "DeiTForImageClassificationWithTeacher" unsupported for now
328}
331class ArchitectureAdapterFactory:
332 """Factory for creating architecture adapters.
334 Supports external registration via `register_adapter()` and automatic
335 discovery of adapters from installed packages via entry points.
336 """
338 _adapters = dict(SUPPORTED_ARCHITECTURES)
339 _entry_points_discovered = False
341 @classmethod
342 def register_adapter(
343 cls, architecture_name: str, adapter_class: type["ArchitectureAdapter"]
344 ) -> None:
345 """Register a custom architecture adapter at runtime.
347 This allows users to add their own architecture adapters without
348 modifying TransformerLens source code.
350 Args:
351 architecture_name: The HuggingFace architecture class name
352 (e.g. ``"Qwen3ForCausalLM"``).
353 adapter_class: The adapter class to register.
355 Example:
356 >>> from transformer_lens.config import TransformerBridgeConfig
357 >>> from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
358 >>> from transformer_lens.factories.architecture_adapter_factory import ArchitectureAdapterFactory
359 >>> class MyAdapter(ArchitectureAdapter):
360 ... def __init__(self, cfg):
361 ... super().__init__(cfg)
362 >>> ArchitectureAdapterFactory.register_adapter("MyModelForCausalLM", MyAdapter)
363 >>> cfg = TransformerBridgeConfig(
364 ... d_model=512, d_head=64, n_layers=6, n_ctx=1024,
365 ... architecture="MyModelForCausalLM",
366 ... )
367 >>> adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
368 >>> isinstance(adapter, MyAdapter)
369 True
370 """
371 cls._adapters[architecture_name] = adapter_class
373 @classmethod
374 def discover_entry_points(cls) -> None:
375 """Discover and register architecture adapters from installed packages.
377 Packages can declare adapters in their ``pyproject.toml``:
378 ```toml
379 [project.entry-points."transformer_lens.architectures"]
380 "MyModelForCausalLM" = "my_package.adapters:MyArchitectureAdapter"
381 ```
382 """
383 if cls._entry_points_discovered:
384 return
385 try:
386 eps = entry_points(group="transformer_lens.architectures")
387 except Exception as e:
388 warnings.warn(
389 f"Failed to discover entry points: {e}. " f"External adapters may not be available."
390 )
391 else:
392 for ep in eps:
393 try:
394 if ep.name in cls._adapters:
395 dist_name = (
396 getattr(ep.dist, "name", "unknown")
397 if ep.dist is not None
398 else "unknown"
399 )
400 warnings.warn(
401 f"Custom architecture adapter {ep.name} provided by {dist_name} "
402 f"attempted to override a native adapter. If you'd like to use this "
403 f"custom adapter, register it explicitly with register_adapter"
404 )
405 continue
406 cls._adapters[ep.name] = ep.load()
407 except Exception as e:
408 warnings.warn(
409 f"Failed to load entry point '{ep.name}': {e}. " f"Skipping this adapter."
410 )
411 cls._entry_points_discovered = True
413 @classmethod
414 def select_architecture_adapter(cls, cfg: TransformerBridgeConfig) -> ArchitectureAdapter:
415 """Select the appropriate architecture adapter for the given config.
417 Args:
418 cfg: The TransformerBridgeConfig to select the adapter for.
420 Returns:
421 The selected architecture adapter.
423 Raises:
424 ValueError: If no adapter is found for the given config.
425 """
426 cls.discover_entry_points()
427 if cfg.architecture is not None:
428 if cfg.architecture in cls._adapters:
429 return cls._adapters[cfg.architecture](cfg)
430 else:
431 raise ValueError(f"Unsupported architecture: {cfg.architecture}")
433 raise ValueError(f"TransformerBridgeConfig must have architecture set, got: {cfg}")