Coverage for transformer_lens/tools/model_registry/generate_report.py: 8%
113 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1#!/usr/bin/env python3
2"""Generate a markdown report of supported and unsupported models.
4This script generates a comprehensive report showing:
5- All supported model IDs grouped by architecture
6- Total count of supported models
7- Unsupported architectures with model counts and descriptions
9Usage:
10 python -m transformer_lens.tools.model_registry.generate_report
11 python -m transformer_lens.tools.model_registry.generate_report --output custom_report.md
12 python -m transformer_lens.tools.model_registry.generate_report --help
13"""
15import argparse
16from datetime import datetime
17from pathlib import Path
19from .api import (
20 get_registry_stats,
21 get_supported_architectures,
22 get_supported_models,
23 get_unsupported_architectures,
24)
26# Descriptions of common architectures (both supported and unsupported)
27ARCHITECTURE_DESCRIPTIONS: dict[str, str] = {
28 # Supported architectures
29 "AfmoeForCausalLM": "Arcee AI's Trinity MoE with sandwich norms, gated attention, and token-choice routing",
30 "ArceeForCausalLM": "Arcee AI's AFM, a Llama-style decoder with an ungated squared-ReLU (ReLU^2) MLP",
31 "ASTForAudioClassification": "MIT's Audio Spectrogram Transformer model for audio classification",
32 "GPT2LMHeadModel": "OpenAI's GPT-2 decoder-only transformer for causal language modeling",
33 "GPTNeoForCausalLM": "EleutherAI's GPT-Neo, an open-source GPT-3-like model",
34 "GPTNeoXForCausalLM": "EleutherAI's GPT-NeoX architecture used in Pythia models",
35 "GPTJForCausalLM": "EleutherAI's GPT-J 6B parameter model",
36 "LlamaForCausalLM": "Meta's LLaMA architecture, basis for many open models",
37 "Llama4ForCausalLM": "Meta's Llama 4 text decoder with interleaved MoE and NoPE layers",
38 "Llama4ForConditionalGeneration": "Meta's Llama 4 Scout/Maverick vision-language model",
39 "LongT5ForConditionalGeneration": "Google's LongT5 with local or transient-global encoder attention",
40 "LEDForConditionalGeneration": "AllenAI's Longformer Encoder-Decoder for long documents",
41 "Florence2ForConditionalGeneration": "Microsoft's Florence-2 vision-language model with DaViT backbone and BART text stack",
42 "Lfm2ForCausalLM": "LiquidAI's LFM2 hybrid convolution/attention dense model",
43 "LLaDAModelLM": "GSAI's bidirectional masked-diffusion language model (single-forward support)",
44 "Lfm2MoeForCausalLM": "LiquidAI's LFM2 hybrid convolution/attention Mixture of Experts model",
45 "MistralForCausalLM": "Mistral AI's efficient 7B parameter model with sliding window attention",
46 "Mistral3ForConditionalGeneration": "Mistral AI's Mistral-Small VLM (Pixtral tower + Mistral decoder)",
47 "MixtralForCausalLM": "Mistral AI's Mixture of Experts model",
48 "GemmaForCausalLM": "Google's Gemma lightweight open model family",
49 "Gemma2ForCausalLM": "Google's Gemma 2 with improved architecture",
50 "Gemma3ForCausalLM": "Google's Gemma 3 latest generation",
51 "Gemma3nForConditionalGeneration": "Google's Gemma 3n efficient tri-modal model (text-only support)",
52 "Gemma4ForConditionalGeneration": "Google's Gemma 4 multimodal model family (text-only support)",
53 "Gemma4UnifiedForConditionalGeneration": "Google's Gemma 4 unified encoder-free multimodal model (text-only support)",
54 "Gemma4ForCausalLM": "Google's Gemma 4 bare text decoder",
55 "DeepseekV4ForCausalLM": "DeepSeek V4 with mHC residual streams and hybrid compressed attention",
56 "GlmMoeDsaForCausalLM": "Z.ai's GLM-5 MoE model with DeepSeek Sparse Attention",
57 "Glm4MoeForCausalLM": "Z.ai's GLM-4.5/4.6/4.7 sparse Mixture-of-Experts causal LM",
58 "Glm4MoeLiteForCausalLM": "Z.ai's GLM-4.7-Flash MoE with DeepSeek-style latent attention",
59 "Qwen2ForCausalLM": "Alibaba's Qwen2 multilingual model",
60 "Qwen2MoeForCausalLM": "Alibaba's Qwen2 sparse Mixture-of-Experts model with shared experts",
61 "Qwen2_5_VLForConditionalGeneration": "Alibaba's Qwen2.5-VL vision-language model (windowed ViT + mRoPE decoder)",
62 "Qwen3VLForConditionalGeneration": "Alibaba's Qwen3-VL with DeepStack multi-level visual injection",
63 "Qwen3VLMoeForConditionalGeneration": "Alibaba's Qwen3-VL-MoE (DeepStack + sparse expert routing)",
64 "Qwen3ForCausalLM": "Alibaba's Qwen3 latest generation",
65 "Qwen3_5ForConditionalGeneration": "Alibaba's Qwen3.5 vision-language model",
66 "Qwen3_5MoeForCausalLM": "Alibaba's Qwen3.5-MoE hybrid linear-attention sparse-MoE text model",
67 "Qwen3_5MoeForConditionalGeneration": "Alibaba's Qwen3.5-MoE vision-language model",
68 "BloomForCausalLM": "BigScience's BLOOM multilingual model",
69 "OPTForCausalLM": "Meta's Open Pre-trained Transformer",
70 "PhiForCausalLM": "Microsoft's Phi small language model",
71 "Phi3ForCausalLM": "Microsoft's Phi-3 improved small model",
72 "PhiMoEForCausalLM": "Microsoft's Phi sparse Mixture of Experts model",
73 "FalconForCausalLM": "TII's Falcon model series",
74 "FalconH1ForCausalLM": "TII's Falcon-H1 parallel attention + Mamba-2 hybrid series",
75 "JambaForCausalLM": "AI21's Jamba interleaved attention + Mamba-1 hybrid with optional MoE",
76 "OlmoForCausalLM": "Allen AI's OLMo open language model",
77 "Olmo2ForCausalLM": "Allen AI's OLMo 2 with improved training",
78 "Olmo3ForCausalLM": "Allen AI's OLMo 3 latest generation",
79 "OlmoeForCausalLM": "Allen AI's OLMoE Mixture of Experts model",
80 "OlmoHybridForCausalLM": "Allen AI's OLMo Hybrid (GatedDeltaNet linear attention + full attention)",
81 "StableLmForCausalLM": "Stability AI's StableLM model",
82 "SmolLM3ForCausalLM": "Hugging Face's SmolLM3 compact open model with NoPE layers",
83 "Starcoder2ForCausalLM": "BigCode's StarCoder2 code generation model",
84 "T5ForConditionalGeneration": "Google's T5 encoder-decoder model (partial support)",
85 "T5WithLMHeadModel": "Legacy T5 class name on old google-t5 checkpoints (t5-3b, t5-11b)",
86 "T5GemmaForConditionalGeneration": "Google's T5Gemma encoder-decoder model with Gemma-style RoPE, GQA, and gated MLP",
87 "BambaForCausalLM": "IBM's Bamba hybrid Mamba-2 + attention decoder",
88 "BartForConditionalGeneration": "Facebook's BART encoder-decoder model",
89 "M2M100ForConditionalGeneration": "Meta's M2M100 / NLLB-200 multilingual translation encoder-decoder",
90 "MarianMTModel": "Helsinki-NLP's Marian opus-mt translation encoder-decoder",
91 "OpenAIGPTLMHeadModel": "The original OpenAI GPT-1 decoder (post-norm, no final LN)",
92 "SeedOssForCausalLM": "ByteDance's Seed-OSS Llama-layout decoder with config-gated biases",
93 "NemotronForCausalLM": "NVIDIA's dense Nemotron/Minitron decoder (LayerNorm1P, squared-ReLU MLP)",
94 "BloomModel": "Headless Bloom checkpoints (loads as BloomForCausalLM with tied embeddings)",
95 "BertLMHeadModel": "BERT with a causal LM head (decoder-style BERT; same module tree as BertForMaskedLM)",
96 "Idefics3ForConditionalGeneration": "HuggingFace's Idefics3/SmolVLM vision-language model (granite-docling)",
97 "Qwen2AudioForConditionalGeneration": "Alibaba's Qwen2-Audio (Whisper-style encoder + Qwen2 decoder)",
98 "Ernie4_5ForCausalLM": "Baidu's dense ERNIE 4.5 Llama-layout decoder",
99 "Ernie4_5_MoeForCausalLM": "Baidu's ERNIE 4.5 MoE (shared experts, dense prefix, interleaved RoPE)",
100 "GlmForCausalLM": "Z.ai's dense GLM-4 decoder (interleaved partial RoPE, joint gate_up MLP)",
101 "Glm4ForCausalLM": "Z.ai's GLM-4-0414 / GLM-Z1 decoder with sandwich norms",
102 "Glm4vForConditionalGeneration": "Z.ai's GLM-4V / GLM-4.1V reasoning vision-language model",
103 "GlmAsrForConditionalGeneration": "Z.ai's GLM-ASR speech recognizer (audio encoder + Llama decoder)",
104 "Starcoder2ForCausalLM": "BigCode's StarCoder2 code model (LayerNorm, biased projections)",
105 "BitNetForCausalLM": "Microsoft's BitNet b1.58 (llama layout with attn/ffn sub-layer norms)",
106 "Exaone4ForCausalLM": "LG AI Research's EXAONE 4.0 (post-norms in residual, per-head QK-norm, hybrid sliding attention)",
107 "BlenderbotForConditionalGeneration": "Meta's Blenderbot dialogue encoder-decoder (asymmetric stacks)",
108 "MiniMaxM2ForCausalLM": "MiniMax's M2 sparse-MoE decoder with sigmoid routing",
109 "ExaoneForCausalLM": "LG AI Research's EXAONE-3.x Korean-English decoder (remote code)",
110 "FalconMambaForCausalLM": "TII's FalconMamba pure-SSM decoder (Mamba-1 with B/C/dt RMS)",
111 "T5Gemma2ForConditionalGeneration": "Google's T5Gemma2 multimodal encoder-decoder model with merged self+cross decoder attention, QK-norm, and dual RoPE (text-only bridge support)",
112 "HrmTextForCausalLM": "Sapient Intelligence's HRM-Text hierarchical two-timescale recurrent model",
113 "BD3LM": "Kuleshov Group's Block Diffusion Language Model (ICLR 2025) for masked text generation",
114 "DreamModel": "HKU-NLP's Dream 7B discrete-diffusion LM (Qwen2.5-based, bidirectional; also Apple DiffuCoder)",
115 "AudioFlamingo3ForConditionalGeneration": "NVIDIA Audio Flamingo 3 audio LALM (Whisper tower + Qwen2.5 LM)",
116 "Emu3ForConditionalGeneration": "BAAI Emu3 unified next-token text+image generation (VQ tokens in a shared vocab)",
117 "FlexOlmoForCausalLM": "AllenAI FlexOlmo federated MoE (merged OLMo-2 experts, inference-time data opt-out)",
118 "MusicFlamingoForConditionalGeneration": "NVIDIA Music Flamingo / Audio Flamingo Next (AF3 + temporal rotary conditioning)",
119 "GiddForDiffusionLM": "GIDD uniform-noise diffusion LM with self-correction (only open non-masked dLM at scale)",
120 "HyenaDNAForCausalLM": "HazyResearch HyenaDNA genomic LM (attention-free Hyena long-conv operator)",
121 "LLaDA2MoeModelLM": "Ant Group LLaDA 2.x masked block-diffusion LM on a DeepSeek-V3-style MoE decoder",
122 "JetMoeForCausalLM": "MIT-IBM JetMoE, the only open at-scale Mixture-of-Attention-heads model",
123 "LagunaForCausalLM": "poolside Laguna coding MoE (heterogeneous per-layer head counts, softplus attention gating)",
124 "ModernBertDecoderForCausalLM": "JHU Ettin decoders (ModernBERT recipe run causally; paired encoder/decoder science suite)",
125 "Jais2ForCausalLM": "G42/Inception Jais 2 Arabic-English family (Nemotron-shaped LN decoder)",
126 "Ministral3ForCausalLM": "Mistral AI Ministral 3 edge models (pure Mistral shape)",
127 "VaultGemmaForCausalLM": "Google VaultGemma, the only fully DP-SGD-pretrained open LLM (Gemma 2 minus post-norms)",
128 "YoutuForCausalLM": "Tencent Youtu-LLM: laptop-scale dense Multi-head Latent Attention (DeepSeek-V2 MLA, all layers dense)",
129 "SwitchTransformersForConditionalGeneration": "Google Switch Transformer, the foundational top-1 capacity-routed MoE (encoder-decoder)",
130 "RwkvForCausalLM": "BlinkDL RWKV-4 WKV linear-attention RNN (Pythia-parallel Pile suite)",
131 "NanoChatForCausalLM": "Karpathy's nanochat education-stack decoder (weightless norms, relu^2 MLP, capped logits)",
132 "HunYuanDenseV1ForCausalLM": "Tencent's open source decoder models",
133 "ViTModel": "Vision Transformer (bare encoder, no classification head)",
134 "ViTForImageClassification": "Vision Transformer with an image classification head",
135 "DeiTModel": "Data-efficient Image Transformer (bare encoder)",
136 "DeiTForImageClassification": "DeiT with a single CLS-token classification head",
137 "Wav2Vec2Model": "Facebook's Wav2Vec 2.0 for speech",
138 "HubertModel": "Facebook's HuBERT for speech",
139 "Cohere2ForCausalLM": "Cohere's Command-A architecture with interleaved sliding-window RoPE and full-attention NoPE layers",
140 "OuroForCausalLM": "ByteDance's Ouro looped language model (LoopLM) with weight-shared iterated depth",
141 "RavenForCausalLM": "tomg-group-umd's Huginn depth-recurrent decoder (prelude / weight-tied recurrent core / coda) with runtime recurrence count",
142 "RWKV7ForCausalLM": 'fla-hub\'s RWKV-7 ("Goose") attention-free recurrent decoder with generalized-delta-rule time-mixing and squared-ReLU channel-mixing',
143 # Unsupported architectures
144 "BertModel": "Google's BERT bidirectional encoder for understanding tasks",
145 "BertForMaskedLM": "BERT with masked language modeling head",
146 "BertForSequenceClassification": "BERT fine-tuned for classification",
147 "RobertaModel": "Facebook's RoBERTa, optimized BERT training",
148 "RobertaForMaskedLM": "RoBERTa with masked language modeling head",
149 "DistilBertModel": "Distilled version of BERT, 40% smaller",
150 "AlbertModel": "A Lite BERT with parameter sharing",
151 "XLNetLMHeadModel": "Google/CMU's XLNet with permutation language modeling",
152 "ElectraModel": "Google's ELECTRA with replaced token detection",
153 "DebertaModel": "Microsoft's DeBERTa with disentangled attention",
154 "DebertaV2Model": "DeBERTa version 2 with improved architecture",
155 "MPNetModel": "Microsoft's MPNet combining MLM and PLM",
156 "LongformerModel": "Allen AI's Longformer for long documents",
157 "BigBirdModel": "Google's BigBird with sparse attention",
158 "ReformerModel": "Google's Reformer with locality-sensitive hashing",
159 "MBartForConditionalGeneration": "Multilingual BART",
160 "PegasusForConditionalGeneration": "Google's PEGASUS for summarization",
161 "MT5ForConditionalGeneration": "Multilingual T5",
162 "WhisperForConditionalGeneration": "OpenAI's Whisper speech recognition",
163 "CLIPModel": "OpenAI's CLIP vision-language model",
164 "SwinModel": "Microsoft's Swin Transformer for vision",
165 "BeitModel": "Microsoft's BERT pre-training for images",
166 "ConvNextModel": "Facebook's ConvNeXt modernized ConvNet",
167 "SegformerModel": "NVIDIA's SegFormer for segmentation",
168 "SpeechT5Model": "Microsoft's SpeechT5 for speech tasks",
169 "BlipModel": "Salesforce's BLIP vision-language model",
170 "Blip2Model": "Salesforce's BLIP-2 with frozen LLM",
171 "LlavaForConditionalGeneration": "Visual instruction-tuned LLaMA",
172 "GitModel": "Microsoft's GIT for vision-language",
173 "PaliGemmaForConditionalGeneration": "Google's PaliGemma vision-language",
174 "CohereForCausalLM": "Cohere's Command models",
175 "DeepseekForCausalLM": "DeepSeek's open models",
176 "InternLMForCausalLM": "Shanghai AI Lab's InternLM",
177 "BaichuanForCausalLM": "Baichuan's Chinese-focused models",
178 "YiForCausalLM": "01.AI's Yi model series",
179 "OrionForCausalLM": "OrionStar's Orion models",
180 "StarcoderForCausalLM": "BigCode's StarCoder for code",
181 "CodeLlamaForCausalLM": "Meta's Code Llama for programming",
182 "CodeGenForCausalLM": "Salesforce's CodeGen models",
183 "SantacoderForCausalLM": "BigCode's SantaCoder",
184}
187def get_architecture_description(arch_id: str) -> str:
188 """Get a description for an architecture, with fallback."""
189 if arch_id in ARCHITECTURE_DESCRIPTIONS:
190 return ARCHITECTURE_DESCRIPTIONS[arch_id]
192 # Generate a basic description from the name
193 if "ForCausalLM" in arch_id:
194 base = arch_id.replace("ForCausalLM", "")
195 return f"{base} architecture for causal language modeling"
196 elif "ForConditionalGeneration" in arch_id:
197 base = arch_id.replace("ForConditionalGeneration", "")
198 return f"{base} encoder-decoder for conditional generation"
199 elif "ForMaskedLM" in arch_id:
200 base = arch_id.replace("ForMaskedLM", "")
201 return f"{base} with masked language modeling head"
202 elif "ForSequenceClassification" in arch_id:
203 base = arch_id.replace("ForSequenceClassification", "")
204 return f"{base} fine-tuned for sequence classification"
205 elif "Model" in arch_id:
206 base = arch_id.replace("Model", "")
207 return f"{base} base model architecture"
208 else:
209 return "Transformer architecture"
212def generate_report(output_path: Path | None = None) -> str:
213 """Generate the markdown report.
215 Args:
216 output_path: Optional path to write the report. If None, only returns the string.
218 Returns:
219 The generated markdown report as a string.
220 """
221 # Gather data
222 models = get_supported_models()
223 architectures = get_supported_architectures()
224 gaps = get_unsupported_architectures()
225 stats = get_registry_stats()
227 # Group models by architecture
228 models_by_arch: dict[str, list[str]] = {}
229 for model in models:
230 arch = model.architecture_id
231 if arch not in models_by_arch:
232 models_by_arch[arch] = []
233 models_by_arch[arch].append(model.model_id)
235 # Sort models within each architecture
236 for arch in models_by_arch:
237 models_by_arch[arch].sort()
239 # Calculate totals
240 total_supported = len(models)
241 total_unsupported = sum(g.total_models for g in gaps)
242 total_all = total_supported + total_unsupported
244 # Build report
245 lines = []
246 lines.append("# TransformerLens Model Compatibility Report")
247 lines.append("")
248 lines.append(f"*Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*")
249 lines.append("")
251 # Summary
252 lines.append("## Summary")
253 lines.append("")
254 lines.append(f"| Metric | Count |")
255 lines.append(f"|--------|-------|")
256 lines.append(f"| Supported Models | {total_supported:,} |")
257 lines.append(f"| Supported Architectures | {len(architectures)} |")
258 lines.append(f"| Verified Models | {stats['total_verified']} |")
259 lines.append(f"| Provisional Models | {stats.get('total_provisional', 0)} |")
260 lines.append(f"| Unsupported Architectures | {len(gaps)} |")
261 lines.append(f"| Models in Unsupported Architectures | {total_unsupported:,} |")
262 lines.append(f"| **Total Potential Models** | **{total_all:,}** |")
263 lines.append("")
265 # Supported models section
266 lines.append("## Supported Models")
267 lines.append("")
268 lines.append(
269 f"TransformerLens supports **{total_supported:,} models** across **{len(architectures)} architectures**."
270 )
271 lines.append("")
273 for arch in sorted(models_by_arch.keys()):
274 model_list = models_by_arch[arch]
275 desc = get_architecture_description(arch)
276 lines.append(f"### {arch}")
277 lines.append("")
278 lines.append(f"*{desc}*")
279 lines.append("")
280 lines.append(f"**{len(model_list)} models:**")
281 lines.append("")
282 for model_id in model_list:
283 model_entry = next((m for m in models if m.model_id == model_id), None)
284 if model_entry and model_entry.status == 1:
285 badge = " ✓"
286 elif model_entry and model_entry.status == 4:
287 badge = " ⚠"
288 else:
289 badge = ""
290 lines.append(f"- `{model_id}`{badge}")
291 lines.append("")
293 # Unsupported architectures section
294 lines.append("## Unsupported Architectures")
295 lines.append("")
296 lines.append(
297 f"The following **{len(gaps)} architectures** are not yet supported by TransformerLens,"
298 )
299 lines.append(f"representing **{total_unsupported:,} models** on HuggingFace.")
300 lines.append("")
301 lines.append("| Architecture | Models | Description |")
302 lines.append("|--------------|--------|-------------|")
304 for gap in gaps:
305 desc = get_architecture_description(gap.architecture_id)
306 lines.append(f"| `{gap.architecture_id}` | {gap.total_models:,} | {desc} |")
308 lines.append("")
310 # Footer
311 lines.append("---")
312 lines.append("")
313 lines.append(
314 "*Report generated by `python -m transformer_lens.tools.model_registry.generate_report`*"
315 )
316 lines.append("")
317 lines.append("✓ = Verified to work with TransformerLens")
318 lines.append(
319 "⚠ = Provisional (structural check only; not numerically verified against HuggingFace)"
320 )
322 report = "\n".join(lines)
324 # Write to file if path provided
325 if output_path:
326 output_path.write_text(report)
327 print(f"Report written to: {output_path}")
329 return report
332def main():
333 """CLI entry point."""
334 parser = argparse.ArgumentParser(
335 description="Generate a markdown report of TransformerLens model compatibility.",
336 formatter_class=argparse.RawDescriptionHelpFormatter,
337 epilog="""
338Examples:
339 # Generate report to default location (MODEL_COMPATIBILITY_REPORT.md)
340 python -m transformer_lens.tools.model_registry.generate_report
342 # Generate report to custom location
343 python -m transformer_lens.tools.model_registry.generate_report -o my_report.md
345 # Print report to stdout only
346 python -m transformer_lens.tools.model_registry.generate_report --stdout
347""",
348 )
349 parser.add_argument(
350 "-o",
351 "--output",
352 type=Path,
353 default=None,
354 help="Output file path (default: MODEL_COMPATIBILITY_REPORT.md in current directory)",
355 )
356 parser.add_argument(
357 "--stdout",
358 action="store_true",
359 help="Print report to stdout instead of writing to file",
360 )
362 args = parser.parse_args()
364 if args.stdout:
365 report = generate_report()
366 print(report)
367 else:
368 output_path = args.output or Path("MODEL_COMPATIBILITY_REPORT.md")
369 generate_report(output_path)
372if __name__ == "__main__": 372 ↛ 373line 372 didn't jump to line 373 because the condition on line 372 was never true
373 main()