Coverage for transformer_lens/tools/model_registry/__init__.py: 100%
8 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Model Registry tools for TransformerLens.
3This package provides tools for discovering and documenting HuggingFace models
4that are compatible with TransformerLens.
6Main modules:
7 - api: Public API for programmatic access to model registry data
8 - schemas: Data classes for model entries, architecture gaps, etc.
9 - verification: Verification tracking for model compatibility
10 - exceptions: Custom exceptions for the model registry
12Example usage:
13 >>> from transformer_lens.tools.model_registry import api # doctest: +SKIP
14 >>> api.is_model_supported("openai-community/gpt2") # doctest: +SKIP
15 True
16 >>> models = api.get_architecture_models("GPT2LMHeadModel") # doctest: +SKIP
17"""
19from .exceptions import (
20 ArchitectureNotSupportedError,
21 DataNotLoadedError,
22 DataValidationError,
23 ModelNotFoundError,
24 ModelRegistryError,
25)
26from .schemas import (
27 ArchitectureAnalysis,
28 ArchitectureGap,
29 ArchitectureGapsReport,
30 ArchitectureStats,
31 ModelEntry,
32 ModelMetadata,
33 ScanInfo,
34 SupportedModelsReport,
35)
36from .verification import VerificationHistory, VerificationRecord
38# Canonical set of HuggingFace architecture class names supported by TransformerBridge.
39# These must match the exact strings found in HF model config.architectures[]
40# and correspond to adapters registered in architecture_adapter_factory.py.
41#
42# Internal-only architectures (NanoGPT, MinGPT, NeelSoluOld, GPT2LMHeadCustomModel,
43# TransformerLensNative) are excluded since they never appear on HuggingFace Hub.
44# Factory-internal alias casings (Gemma1, Neo, NeoX) are also excluded since they
45# route to canonical adapters but HF reports the canonical names (Gemma, GPTNeo,
46# GPTNeoX) in config.architectures instead.
47HF_SUPPORTED_ARCHITECTURES: set[str] = {
48 "AfmoeForCausalLM",
49 "ApertusForCausalLM",
50 "ArceeForCausalLM",
51 "ASTForAudioClassification",
52 "BaiChuanForCausalLM",
53 "BaichuanForCausalLM",
54 "BambaForCausalLM",
55 "BartForConditionalGeneration",
56 "BD3LM",
57 "DreamModel",
58 "Emu3ForConditionalGeneration",
59 "AudioFlamingo3ForConditionalGeneration",
60 "FlexOlmoForCausalLM",
61 "GiddForDiffusionLM",
62 "HyenaDNAForCausalLM",
63 "LLaDA2MoeModelLM",
64 "Jais2ForCausalLM",
65 "JetMoeForCausalLM",
66 "JetMoEForCausalLM",
67 "LagunaForCausalLM",
68 "Ministral3ForCausalLM",
69 "VaultGemmaForCausalLM",
70 "YoutuForCausalLM",
71 "ModernBertDecoderForCausalLM",
72 "MusicFlamingoForConditionalGeneration",
73 "NanoChatForCausalLM",
74 "RwkvForCausalLM",
75 "SwitchTransformersForConditionalGeneration",
76 "BertForMaskedLM",
77 "BertLMHeadModel",
78 "BitNetForCausalLM",
79 "BlenderbotForConditionalGeneration",
80 "BloomForCausalLM",
81 "BloomModel",
82 "CodeGenForCausalLM",
83 "Cohere2ForCausalLM",
84 "CohereForCausalLM",
85 "DeepseekV2ForCausalLM",
86 "DeepseekV3ForCausalLM",
87 "Ernie4_5ForCausalLM",
88 "Ernie4_5_MoeForCausalLM",
89 "ExaoneForCausalLM",
90 "Exaone4ForCausalLM",
91 "DeepseekV4ForCausalLM",
92 "FalconForCausalLM",
93 "FalconH1ForCausalLM",
94 "FalconMambaForCausalLM",
95 "Florence2ForConditionalGeneration",
96 "GemmaForCausalLM",
97 "Gemma2ForCausalLM",
98 "Gemma3ForCausalLM",
99 "Gemma3ForConditionalGeneration",
100 "Gemma3nForConditionalGeneration",
101 "Gemma4ForConditionalGeneration",
102 "Gemma4UnifiedForConditionalGeneration",
103 "Gemma4ForCausalLM",
104 "GlmForCausalLM",
105 "Glm4ForCausalLM",
106 "Glm4vForConditionalGeneration",
107 "GlmAsrForConditionalGeneration",
108 "Glm4MoeForCausalLM",
109 "Glm4MoeLiteForCausalLM",
110 "GlmMoeDsaForCausalLM",
111 "GraniteForCausalLM",
112 "GraniteMoeForCausalLM",
113 "GraniteMoeHybridForCausalLM",
114 "GPT2LMHeadModel",
115 "GPTBigCodeForCausalLM",
116 "GptOssForCausalLM",
117 "GPTJForCausalLM",
118 "GPTNeoForCausalLM",
119 "OpenAIGPTLMHeadModel",
120 "OpenELMForCausalLM",
121 "GPTNeoXForCausalLM",
122 "HubertForCTC",
123 "HubertModel",
124 "Wav2Vec2ForCTC",
125 "Wav2Vec2ForPreTraining",
126 "Wav2Vec2Model",
127 "HrmTextForCausalLM",
128 "HunYuanDenseV1ForCausalLM",
129 "Idefics3ForConditionalGeneration",
130 "InternLM2ForCausalLM",
131 "JambaForCausalLM",
132 "LEDForConditionalGeneration",
133 "LLaDAModelLM",
134 "LlamaForCausalLM",
135 "Llama4ForCausalLM",
136 "Llama4ForConditionalGeneration",
137 "LlavaForConditionalGeneration",
138 "LlavaNextForConditionalGeneration",
139 "LlavaOnevisionForConditionalGeneration",
140 "Lfm2ForCausalLM",
141 "Lfm2MoeForCausalLM",
142 "LongT5ForConditionalGeneration",
143 "M2M100ForConditionalGeneration",
144 "MambaForCausalLM",
145 "Mamba2ForCausalLM",
146 "MarianMTModel",
147 "MBartForConditionalGeneration",
148 "MiniMaxM2ForCausalLM",
149 "NemotronForCausalLM",
150 "NemotronHForCausalLM",
151 "MPTForCausalLM",
152 "MptForCausalLM",
153 "MistralForCausalLM",
154 "Mistral3ForConditionalGeneration",
155 "MixtralForCausalLM",
156 "Olmo2ForCausalLM",
157 "Olmo3ForCausalLM",
158 "OlmoForCausalLM",
159 "OlmoeForCausalLM",
160 "OlmoHybridForCausalLM",
161 "OPTForCausalLM",
162 "PegasusForConditionalGeneration",
163 "OuroForCausalLM",
164 "PhiForCausalLM",
165 "Phi3ForCausalLM",
166 "PhiMoEForCausalLM",
167 "QwenForCausalLM",
168 "Qwen2ForCausalLM",
169 "Qwen2_5_VLForConditionalGeneration",
170 "Qwen3VLForConditionalGeneration",
171 "Qwen3VLMoeForConditionalGeneration",
172 "Qwen2AudioForConditionalGeneration",
173 "Qwen2MoeForCausalLM",
174 "Qwen3ForCausalLM",
175 "Qwen3MoeForCausalLM",
176 "Qwen3NextForCausalLM",
177 "Qwen3_5ForCausalLM",
178 "Qwen3_5ForConditionalGeneration",
179 "Qwen3_5MoeForCausalLM",
180 "Qwen3_5MoeForConditionalGeneration",
181 "RavenForCausalLM",
182 "RecurrentGemmaForCausalLM",
183 "SeedOssForCausalLM",
184 "RWKV7ForCausalLM",
185 "SmolLM3ForCausalLM",
186 "StableLmForCausalLM",
187 "Starcoder2ForCausalLM",
188 "T5ForConditionalGeneration",
189 "T5WithLMHeadModel",
190 "MT5ForConditionalGeneration",
191 "T5GemmaForConditionalGeneration",
192 "T5Gemma2ForConditionalGeneration",
193 "XGLMForCausalLM",
194 "ViTModel",
195 "ViTForImageClassification",
196 "DeiTModel",
197 "DeiTForImageClassification",
198 "Zamba2ForCausalLM",
199}
201# Hub checkpoint configs can disagree with the transformers class name on casing
202# (jetmoe/jetmoe-8b reports "JetMoEForCausalLM"; registry rows are keyed by the
203# transformers class "JetMoeForCausalLM"). Row-scanning lookups normalize through
204# this map before matching.
205ARCHITECTURE_ALIASES: dict[str, str] = {
206 "JetMoEForCausalLM": "JetMoeForCausalLM",
207}
209# Foundation-trained orgs per architecture. Source of truth for the scraper's
210# download-threshold bypass and the docs table's "Canonical only" toggle.
211CANONICAL_AUTHORS_BY_ARCH: dict[str, list[str]] = {
212 "AfmoeForCausalLM": ["arcee-ai"],
213 "ApertusForCausalLM": ["swiss-ai"],
214 "ArceeForCausalLM": ["arcee-ai"],
215 "ASTForAudioClassification": ["MIT"],
216 "BaiChuanForCausalLM": ["baichuan-inc"],
217 "BaichuanForCausalLM": ["baichuan-inc"],
218 "BambaForCausalLM": ["ibm-ai-platform"],
219 "BartForConditionalGeneration": ["facebook"],
220 "BD3LM": ["kuleshov-group"],
221 "DreamModel": ["Dream-org", "apple"],
222 "Emu3ForConditionalGeneration": ["BAAI"],
223 "AudioFlamingo3ForConditionalGeneration": ["nvidia"],
224 "FlexOlmoForCausalLM": ["allenai"],
225 "GiddForDiffusionLM": ["dvruette"],
226 "HyenaDNAForCausalLM": ["LongSafari"],
227 "LLaDA2MoeModelLM": ["inclusionAI"],
228 "Jais2ForCausalLM": ["inceptionai"],
229 "JetMoeForCausalLM": ["jetmoe"],
230 "JetMoEForCausalLM": ["jetmoe"],
231 "LagunaForCausalLM": ["poolside"],
232 "Ministral3ForCausalLM": ["mistralai"],
233 "VaultGemmaForCausalLM": ["google"],
234 "YoutuForCausalLM": ["tencent"],
235 "ModernBertDecoderForCausalLM": ["jhu-clsp"],
236 "MusicFlamingoForConditionalGeneration": ["nvidia"],
237 "NanoChatForCausalLM": ["nanochat-students", "karpathy"],
238 "RwkvForCausalLM": ["RWKV"],
239 "SwitchTransformersForConditionalGeneration": ["google"],
240 "BertForMaskedLM": ["google-bert"],
241 "BertLMHeadModel": ["google-bert"],
242 "BitNetForCausalLM": ["microsoft"],
243 "BlenderbotForConditionalGeneration": ["facebook"],
244 "BloomForCausalLM": ["bigscience"],
245 "BloomModel": ["bigscience", "TurkuNLP"],
246 "CodeGenForCausalLM": ["Salesforce"],
247 "Cohere2ForCausalLM": ["CohereLabs"],
248 "CohereForCausalLM": ["CohereLabs"],
249 "DeepseekV2ForCausalLM": ["deepseek-ai"],
250 "DeepseekV3ForCausalLM": ["deepseek-ai"],
251 "Ernie4_5ForCausalLM": ["baidu"],
252 "Ernie4_5_MoeForCausalLM": ["baidu"],
253 "ExaoneForCausalLM": ["LGAI-EXAONE"],
254 "Exaone4ForCausalLM": ["LGAI-EXAONE"],
255 "DeepseekV4ForCausalLM": ["deepseek-ai"],
256 "FalconForCausalLM": ["tiiuae"],
257 "FalconH1ForCausalLM": ["tiiuae"],
258 "FalconMambaForCausalLM": ["tiiuae"],
259 "Florence2ForConditionalGeneration": ["florence-community", "microsoft"],
260 "Gemma2ForCausalLM": ["google"],
261 "Gemma3ForCausalLM": ["google"],
262 "Gemma3ForConditionalGeneration": ["google"],
263 "Gemma3nForConditionalGeneration": ["google"],
264 "Gemma4ForConditionalGeneration": ["google"],
265 "Gemma4UnifiedForConditionalGeneration": ["google"],
266 "Gemma4ForCausalLM": ["google"],
267 "GemmaForCausalLM": ["google"],
268 "GlmForCausalLM": ["zai-org"],
269 "Glm4ForCausalLM": ["zai-org", "THUDM"],
270 "Glm4vForConditionalGeneration": ["zai-org", "THUDM"],
271 "GlmAsrForConditionalGeneration": ["zai-org"],
272 "Glm4MoeForCausalLM": ["zai-org"],
273 "Glm4MoeLiteForCausalLM": ["zai-org"],
274 "GlmMoeDsaForCausalLM": ["zai-org"],
275 "GPT2LMHeadModel": ["openai-community", "stanford-crfm", "Writer"],
276 "GPTBigCodeForCausalLM": ["bigcode"],
277 "GptOssForCausalLM": ["openai"],
278 "GPTJForCausalLM": ["EleutherAI", "togethercomputer"],
279 "GPTNeoForCausalLM": ["EleutherAI", "roneneldan"],
280 "GPTNeoXForCausalLM": ["EleutherAI", "cyberagent", "stabilityai", "togethercomputer"],
281 "GraniteForCausalLM": ["ibm-granite"],
282 "GraniteMoeForCausalLM": ["ibm-granite"],
283 "GraniteMoeHybridForCausalLM": ["ibm-granite"],
284 "HubertForCTC": ["facebook"],
285 "HubertModel": ["facebook"],
286 "Wav2Vec2ForCTC": ["facebook"],
287 "Wav2Vec2ForPreTraining": ["facebook"],
288 "Wav2Vec2Model": ["facebook"],
289 "HrmTextForCausalLM": ["sapientinc"],
290 "HunYuanDenseV1ForCausalLM": ["tencent"],
291 "Idefics3ForConditionalGeneration": ["HuggingFaceM4", "ibm-granite"],
292 "InternLM2ForCausalLM": ["internlm"],
293 "JambaForCausalLM": ["ai21labs"],
294 "LEDForConditionalGeneration": ["allenai"],
295 "LLaDAModelLM": ["GSAI-ML"],
296 "LlamaForCausalLM": ["meta-llama", "huggyllama", "codellama", "SimpleStories"],
297 "Llama4ForCausalLM": ["meta-llama", "trl-internal-testing"],
298 "Llama4ForConditionalGeneration": ["meta-llama"],
299 "LlavaForConditionalGeneration": ["llava-hf"],
300 "LlavaNextForConditionalGeneration": ["llava-hf"],
301 "LlavaOnevisionForConditionalGeneration": ["llava-hf"],
302 "Lfm2ForCausalLM": ["LiquidAI"],
303 "Lfm2MoeForCausalLM": ["LiquidAI"],
304 "Mamba2ForCausalLM": ["state-spaces"],
305 "LongT5ForConditionalGeneration": ["google"],
306 "M2M100ForConditionalGeneration": ["facebook"],
307 "MambaForCausalLM": ["state-spaces"],
308 "MarianMTModel": ["Helsinki-NLP"],
309 "MBartForConditionalGeneration": ["facebook"],
310 "MiniMaxM2ForCausalLM": ["MiniMaxAI"],
311 "NemotronForCausalLM": ["nvidia"],
312 "NemotronHForCausalLM": ["nvidia"],
313 "MistralForCausalLM": ["mistralai"],
314 "Mistral3ForConditionalGeneration": ["mistralai"],
315 "MixtralForCausalLM": ["mistralai"],
316 "MPTForCausalLM": ["mosaicml"],
317 "MptForCausalLM": ["mosaicml"],
318 "MT5ForConditionalGeneration": ["google", "bigscience", "csebuetnlp"],
319 "Olmo2ForCausalLM": ["allenai", "HPLT"],
320 "Olmo3ForCausalLM": ["allenai"],
321 "OlmoeForCausalLM": ["allenai"],
322 "OlmoHybridForCausalLM": ["allenai"],
323 "OlmoForCausalLM": ["allenai"],
324 "OpenAIGPTLMHeadModel": ["openai-community"],
325 "OpenELMForCausalLM": ["apple"],
326 "OPTForCausalLM": ["facebook"],
327 "PegasusForConditionalGeneration": ["google"],
328 "OuroForCausalLM": ["ByteDance"],
329 "Phi3ForCausalLM": ["microsoft"],
330 "PhiMoEForCausalLM": ["microsoft"],
331 "PhiForCausalLM": ["microsoft"],
332 "Qwen2ForCausalLM": ["Qwen", "nvidia"],
333 "Qwen2_5_VLForConditionalGeneration": ["Qwen"],
334 "Qwen3VLForConditionalGeneration": ["Qwen"],
335 "Qwen3VLMoeForConditionalGeneration": ["Qwen"],
336 "Qwen2AudioForConditionalGeneration": ["Qwen"],
337 "Qwen2MoeForCausalLM": ["Qwen"],
338 "Qwen3ForCausalLM": ["Qwen", "nvidia"],
339 "Qwen3MoeForCausalLM": ["Qwen"],
340 "Qwen3NextForCausalLM": ["Qwen"],
341 "Qwen3_5ForCausalLM": ["Qwen"],
342 "Qwen3_5ForConditionalGeneration": ["Qwen"],
343 "Qwen3_5MoeForCausalLM": ["Qwen"],
344 "Qwen3_5MoeForConditionalGeneration": ["Qwen"],
345 "QwenForCausalLM": ["Qwen"],
346 "RavenForCausalLM": ["tomg-group-umd"],
347 "RecurrentGemmaForCausalLM": ["google"],
348 "RWKV7ForCausalLM": ["fla-hub"],
349 "SeedOssForCausalLM": ["ByteDance-Seed"],
350 "SmolLM3ForCausalLM": ["HuggingFaceTB"],
351 "StableLmForCausalLM": ["stabilityai"],
352 "Starcoder2ForCausalLM": ["bigcode"],
353 "T5ForConditionalGeneration": ["google-t5", "google", "Salesforce", "MBZUAI"],
354 "T5WithLMHeadModel": ["google-t5"],
355 "T5GemmaForConditionalGeneration": ["google"],
356 "T5Gemma2ForConditionalGeneration": ["google"],
357 "XGLMForCausalLM": ["facebook"],
358 "ViTModel": ["google", "facebook"],
359 "ViTForImageClassification": ["google", "facebook"],
360 "DeiTModel": ["facebook"],
361 "DeiTForImageClassification": ["facebook"],
362 "Zamba2ForCausalLM": ["Zyphra"],
363}
365# Model-name prefixes that require trust_remote_code=True to load (custom
366# modeling code on the HF Hub).
367REMOTE_CODE_MODEL_PREFIXES: tuple[str, ...] = (
368 "apple/DiffuCoder", # DreamModel (DiffuCoder) — same remote code family
369 "apple/OpenELM",
370 "baichuan-inc/", # BaichuanForCausalLM — ships own modeling_baichuan.py
371 "ByteDance/Ouro-", # OuroForCausalLM — ships own modeling_ouro.py
372 "bigcode/santacoder",
373 "Dream-org/", # DreamModel — ships own modeling_dream.py
374 "dvruette/", # GiddForDiffusionLM — ships own modeling_gidd.py
375 "inclusionAI/", # LLaDA2MoeModelLM — ships own modeling_llada2_moe.py
376 "kuleshov-group/", # BD3LM — ships own custom modeling_d_dit.py
377 "LGAI-EXAONE/", # ExaoneForCausalLM (EXAONE-3.x) — ships own modeling_exaone.py
378 "LongSafari/", # HyenaDNAForCausalLM — ships own modeling_hyena.py
379 "microsoft/phi-2",
380 "microsoft/phi-4",
381 "openai/gpt-oss-",
382 "poolside/", # LagunaForCausalLM — ships own modeling_laguna.py
383 "Qwen/Qwen-",
384 "Qwen/Qwen3-",
385 "swiss-ai/Apertus-",
386 "internlm/", # InternLM2ForCausalLM — ships own modeling_internlm2.py
387 "GSAI-ML/LLaDA", # LLaDAModelLM — ships configuration_llada.py/modeling_llada.py
388)
390__all__ = [
391 # Constants
392 "HF_SUPPORTED_ARCHITECTURES",
393 "ARCHITECTURE_ALIASES",
394 "CANONICAL_AUTHORS_BY_ARCH",
395 "REMOTE_CODE_MODEL_PREFIXES",
396 # Exceptions
397 "ModelRegistryError",
398 "ModelNotFoundError",
399 "ArchitectureNotSupportedError",
400 "DataNotLoadedError",
401 "DataValidationError",
402 # Schemas
403 "ModelEntry",
404 "ModelMetadata",
405 "ScanInfo",
406 "ArchitectureGap",
407 "ArchitectureStats",
408 "ArchitectureAnalysis",
409 "SupportedModelsReport",
410 "ArchitectureGapsReport",
411 # Verification
412 "VerificationRecord",
413 "VerificationHistory",
414]