Coverage for transformer_lens/tools/model_registry/__init__.py: 100%
6 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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 "HrmTextForCausalLM",
125 "HunYuanDenseV1ForCausalLM",
126 "Idefics3ForConditionalGeneration",
127 "InternLM2ForCausalLM",
128 "JambaForCausalLM",
129 "LEDForConditionalGeneration",
130 "LLaDAModelLM",
131 "LlamaForCausalLM",
132 "Llama4ForCausalLM",
133 "Llama4ForConditionalGeneration",
134 "LlavaForConditionalGeneration",
135 "LlavaNextForConditionalGeneration",
136 "LlavaOnevisionForConditionalGeneration",
137 "Lfm2ForCausalLM",
138 "Lfm2MoeForCausalLM",
139 "LongT5ForConditionalGeneration",
140 "M2M100ForConditionalGeneration",
141 "MambaForCausalLM",
142 "Mamba2ForCausalLM",
143 "MarianMTModel",
144 "MBartForConditionalGeneration",
145 "MiniMaxM2ForCausalLM",
146 "NemotronForCausalLM",
147 "NemotronHForCausalLM",
148 "MPTForCausalLM",
149 "MptForCausalLM",
150 "MistralForCausalLM",
151 "Mistral3ForConditionalGeneration",
152 "MixtralForCausalLM",
153 "Olmo2ForCausalLM",
154 "Olmo3ForCausalLM",
155 "OlmoForCausalLM",
156 "OlmoeForCausalLM",
157 "OlmoHybridForCausalLM",
158 "OPTForCausalLM",
159 "PegasusForConditionalGeneration",
160 "OuroForCausalLM",
161 "PhiForCausalLM",
162 "Phi3ForCausalLM",
163 "PhiMoEForCausalLM",
164 "QwenForCausalLM",
165 "Qwen2ForCausalLM",
166 "Qwen2_5_VLForConditionalGeneration",
167 "Qwen3VLForConditionalGeneration",
168 "Qwen3VLMoeForConditionalGeneration",
169 "Qwen2AudioForConditionalGeneration",
170 "Qwen2MoeForCausalLM",
171 "Qwen3ForCausalLM",
172 "Qwen3MoeForCausalLM",
173 "Qwen3NextForCausalLM",
174 "Qwen3_5ForCausalLM",
175 "Qwen3_5ForConditionalGeneration",
176 "Qwen3_5MoeForCausalLM",
177 "Qwen3_5MoeForConditionalGeneration",
178 "RavenForCausalLM",
179 "RecurrentGemmaForCausalLM",
180 "SeedOssForCausalLM",
181 "RWKV7ForCausalLM",
182 "SmolLM3ForCausalLM",
183 "StableLmForCausalLM",
184 "Starcoder2ForCausalLM",
185 "T5ForConditionalGeneration",
186 "T5WithLMHeadModel",
187 "MT5ForConditionalGeneration",
188 "T5GemmaForConditionalGeneration",
189 "T5Gemma2ForConditionalGeneration",
190 "XGLMForCausalLM",
191 "ViTModel",
192 "ViTForImageClassification",
193 "DeiTModel",
194 "DeiTForImageClassification",
195 "Zamba2ForCausalLM",
196}
198# Foundation-trained orgs per architecture. Source of truth for the scraper's
199# download-threshold bypass and the docs table's "Canonical only" toggle.
200CANONICAL_AUTHORS_BY_ARCH: dict[str, list[str]] = {
201 "AfmoeForCausalLM": ["arcee-ai"],
202 "ApertusForCausalLM": ["swiss-ai"],
203 "ArceeForCausalLM": ["arcee-ai"],
204 "ASTForAudioClassification": ["MIT"],
205 "BaiChuanForCausalLM": ["baichuan-inc"],
206 "BaichuanForCausalLM": ["baichuan-inc"],
207 "BambaForCausalLM": ["ibm-ai-platform"],
208 "BartForConditionalGeneration": ["facebook"],
209 "BD3LM": ["kuleshov-group"],
210 "DreamModel": ["Dream-org", "apple"],
211 "Emu3ForConditionalGeneration": ["BAAI"],
212 "AudioFlamingo3ForConditionalGeneration": ["nvidia"],
213 "FlexOlmoForCausalLM": ["allenai"],
214 "GiddForDiffusionLM": ["dvruette"],
215 "HyenaDNAForCausalLM": ["LongSafari"],
216 "LLaDA2MoeModelLM": ["inclusionAI"],
217 "Jais2ForCausalLM": ["inceptionai"],
218 "JetMoeForCausalLM": ["jetmoe"],
219 "JetMoEForCausalLM": ["jetmoe"],
220 "LagunaForCausalLM": ["poolside"],
221 "Ministral3ForCausalLM": ["mistralai"],
222 "VaultGemmaForCausalLM": ["google"],
223 "YoutuForCausalLM": ["tencent"],
224 "ModernBertDecoderForCausalLM": ["jhu-clsp"],
225 "MusicFlamingoForConditionalGeneration": ["nvidia"],
226 "NanoChatForCausalLM": ["nanochat-students", "karpathy"],
227 "RwkvForCausalLM": ["RWKV"],
228 "SwitchTransformersForConditionalGeneration": ["google"],
229 "BertForMaskedLM": ["google-bert"],
230 "BertLMHeadModel": ["google-bert"],
231 "BitNetForCausalLM": ["microsoft"],
232 "BlenderbotForConditionalGeneration": ["facebook"],
233 "BloomForCausalLM": ["bigscience"],
234 "BloomModel": ["bigscience", "TurkuNLP"],
235 "CodeGenForCausalLM": ["Salesforce"],
236 "Cohere2ForCausalLM": ["CohereLabs"],
237 "CohereForCausalLM": ["CohereLabs"],
238 "DeepseekV2ForCausalLM": ["deepseek-ai"],
239 "DeepseekV3ForCausalLM": ["deepseek-ai"],
240 "Ernie4_5ForCausalLM": ["baidu"],
241 "Ernie4_5_MoeForCausalLM": ["baidu"],
242 "ExaoneForCausalLM": ["LGAI-EXAONE"],
243 "Exaone4ForCausalLM": ["LGAI-EXAONE"],
244 "DeepseekV4ForCausalLM": ["deepseek-ai"],
245 "FalconForCausalLM": ["tiiuae"],
246 "FalconH1ForCausalLM": ["tiiuae"],
247 "FalconMambaForCausalLM": ["tiiuae"],
248 "Florence2ForConditionalGeneration": ["florence-community", "microsoft"],
249 "Gemma2ForCausalLM": ["google"],
250 "Gemma3ForCausalLM": ["google"],
251 "Gemma3ForConditionalGeneration": ["google"],
252 "Gemma3nForConditionalGeneration": ["google"],
253 "Gemma4ForConditionalGeneration": ["google"],
254 "Gemma4UnifiedForConditionalGeneration": ["google"],
255 "Gemma4ForCausalLM": ["google"],
256 "GemmaForCausalLM": ["google"],
257 "GlmForCausalLM": ["zai-org"],
258 "Glm4ForCausalLM": ["zai-org", "THUDM"],
259 "Glm4vForConditionalGeneration": ["zai-org", "THUDM"],
260 "GlmAsrForConditionalGeneration": ["zai-org"],
261 "Glm4MoeForCausalLM": ["zai-org"],
262 "Glm4MoeLiteForCausalLM": ["zai-org"],
263 "GlmMoeDsaForCausalLM": ["zai-org"],
264 "GPT2LMHeadModel": ["openai-community", "stanford-crfm", "Writer"],
265 "GPTBigCodeForCausalLM": ["bigcode"],
266 "GptOssForCausalLM": ["openai"],
267 "GPTJForCausalLM": ["EleutherAI", "togethercomputer"],
268 "GPTNeoForCausalLM": ["EleutherAI", "roneneldan"],
269 "GPTNeoXForCausalLM": ["EleutherAI", "cyberagent", "stabilityai", "togethercomputer"],
270 "GraniteForCausalLM": ["ibm-granite"],
271 "GraniteMoeForCausalLM": ["ibm-granite"],
272 "GraniteMoeHybridForCausalLM": ["ibm-granite"],
273 "HubertForCTC": ["facebook"],
274 "HubertModel": ["facebook"],
275 "HrmTextForCausalLM": ["sapientinc"],
276 "HunYuanDenseV1ForCausalLM": ["tencent"],
277 "Idefics3ForConditionalGeneration": ["HuggingFaceM4", "ibm-granite"],
278 "InternLM2ForCausalLM": ["internlm"],
279 "JambaForCausalLM": ["ai21labs"],
280 "LEDForConditionalGeneration": ["allenai"],
281 "LLaDAModelLM": ["GSAI-ML"],
282 "LlamaForCausalLM": ["meta-llama", "huggyllama", "codellama", "SimpleStories"],
283 "Llama4ForCausalLM": ["meta-llama", "trl-internal-testing"],
284 "Llama4ForConditionalGeneration": ["meta-llama"],
285 "LlavaForConditionalGeneration": ["llava-hf"],
286 "LlavaNextForConditionalGeneration": ["llava-hf"],
287 "LlavaOnevisionForConditionalGeneration": ["llava-hf"],
288 "Lfm2ForCausalLM": ["LiquidAI"],
289 "Lfm2MoeForCausalLM": ["LiquidAI"],
290 "Mamba2ForCausalLM": ["state-spaces"],
291 "LongT5ForConditionalGeneration": ["google"],
292 "M2M100ForConditionalGeneration": ["facebook"],
293 "MambaForCausalLM": ["state-spaces"],
294 "MarianMTModel": ["Helsinki-NLP"],
295 "MBartForConditionalGeneration": ["facebook"],
296 "MiniMaxM2ForCausalLM": ["MiniMaxAI"],
297 "NemotronForCausalLM": ["nvidia"],
298 "NemotronHForCausalLM": ["nvidia"],
299 "MistralForCausalLM": ["mistralai"],
300 "Mistral3ForConditionalGeneration": ["mistralai"],
301 "MixtralForCausalLM": ["mistralai"],
302 "MPTForCausalLM": ["mosaicml"],
303 "MptForCausalLM": ["mosaicml"],
304 "MT5ForConditionalGeneration": ["google", "bigscience", "csebuetnlp"],
305 "Olmo2ForCausalLM": ["allenai", "HPLT"],
306 "Olmo3ForCausalLM": ["allenai"],
307 "OlmoeForCausalLM": ["allenai"],
308 "OlmoHybridForCausalLM": ["allenai"],
309 "OlmoForCausalLM": ["allenai"],
310 "OpenAIGPTLMHeadModel": ["openai-community"],
311 "OpenELMForCausalLM": ["apple"],
312 "OPTForCausalLM": ["facebook"],
313 "PegasusForConditionalGeneration": ["google"],
314 "OuroForCausalLM": ["ByteDance"],
315 "Phi3ForCausalLM": ["microsoft"],
316 "PhiMoEForCausalLM": ["microsoft"],
317 "PhiForCausalLM": ["microsoft"],
318 "Qwen2ForCausalLM": ["Qwen", "nvidia"],
319 "Qwen2_5_VLForConditionalGeneration": ["Qwen"],
320 "Qwen3VLForConditionalGeneration": ["Qwen"],
321 "Qwen3VLMoeForConditionalGeneration": ["Qwen"],
322 "Qwen2AudioForConditionalGeneration": ["Qwen"],
323 "Qwen2MoeForCausalLM": ["Qwen"],
324 "Qwen3ForCausalLM": ["Qwen", "nvidia"],
325 "Qwen3MoeForCausalLM": ["Qwen"],
326 "Qwen3NextForCausalLM": ["Qwen"],
327 "Qwen3_5ForCausalLM": ["Qwen"],
328 "Qwen3_5ForConditionalGeneration": ["Qwen"],
329 "Qwen3_5MoeForCausalLM": ["Qwen"],
330 "Qwen3_5MoeForConditionalGeneration": ["Qwen"],
331 "QwenForCausalLM": ["Qwen"],
332 "RavenForCausalLM": ["tomg-group-umd"],
333 "RecurrentGemmaForCausalLM": ["google"],
334 "SeedOssForCausalLM": ["ByteDance-Seed"],
335 "RWKV7ForCausalLM": ["fla-hub"],
336 "SmolLM3ForCausalLM": ["HuggingFaceTB"],
337 "StableLmForCausalLM": ["stabilityai"],
338 "Starcoder2ForCausalLM": ["bigcode"],
339 "T5ForConditionalGeneration": ["google-t5", "google", "Salesforce", "MBZUAI"],
340 "T5WithLMHeadModel": ["google-t5"],
341 "T5GemmaForConditionalGeneration": ["google"],
342 "T5Gemma2ForConditionalGeneration": ["google"],
343 "XGLMForCausalLM": ["facebook"],
344 "ViTModel": ["google"],
345 "ViTForImageClassification": ["google"],
346 "DeiTModel": ["facebook"],
347 "DeiTForImageClassification": ["facebook"],
348 "Zamba2ForCausalLM": ["Zyphra"],
349}
351__all__ = [
352 # Constants
353 "HF_SUPPORTED_ARCHITECTURES",
354 "CANONICAL_AUTHORS_BY_ARCH",
355 # Exceptions
356 "ModelRegistryError",
357 "ModelNotFoundError",
358 "ArchitectureNotSupportedError",
359 "DataNotLoadedError",
360 "DataValidationError",
361 # Schemas
362 "ModelEntry",
363 "ModelMetadata",
364 "ScanInfo",
365 "ArchitectureGap",
366 "ArchitectureStats",
367 "ArchitectureAnalysis",
368 "SupportedModelsReport",
369 "ArchitectureGapsReport",
370 # Verification
371 "VerificationRecord",
372 "VerificationHistory",
373]