Coverage for transformer_lens/tools/model_registry/generate_report.py: 7%
113 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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 "T5ForConditionalGeneration": "Google's T5 encoder-decoder model (partial support)",
84 "T5WithLMHeadModel": "Legacy T5 class name on old google-t5 checkpoints (t5-3b, t5-11b)",
85 "T5GemmaForConditionalGeneration": "Google's T5Gemma encoder-decoder model with Gemma-style RoPE, GQA, and gated MLP",
86 "BambaForCausalLM": "IBM's Bamba hybrid Mamba-2 + attention decoder",
87 "BartForConditionalGeneration": "Facebook's BART encoder-decoder model",
88 "M2M100ForConditionalGeneration": "Meta's M2M100 / NLLB-200 multilingual translation encoder-decoder",
89 "MarianMTModel": "Helsinki-NLP's Marian opus-mt translation encoder-decoder",
90 "OpenAIGPTLMHeadModel": "The original OpenAI GPT-1 decoder (post-norm, no final LN)",
91 "SeedOssForCausalLM": "ByteDance's Seed-OSS Llama-layout decoder with config-gated biases",
92 "NemotronForCausalLM": "NVIDIA's dense Nemotron/Minitron decoder (LayerNorm1P, squared-ReLU MLP)",
93 "BloomModel": "Headless Bloom checkpoints (loads as BloomForCausalLM with tied embeddings)",
94 "BertLMHeadModel": "BERT with a causal LM head (decoder-style BERT; same module tree as BertForMaskedLM)",
95 "Idefics3ForConditionalGeneration": "HuggingFace's Idefics3/SmolVLM vision-language model (granite-docling)",
96 "Qwen2AudioForConditionalGeneration": "Alibaba's Qwen2-Audio (Whisper-style encoder + Qwen2 decoder)",
97 "Ernie4_5ForCausalLM": "Baidu's dense ERNIE 4.5 Llama-layout decoder",
98 "Ernie4_5_MoeForCausalLM": "Baidu's ERNIE 4.5 MoE (shared experts, dense prefix, interleaved RoPE)",
99 "GlmForCausalLM": "Z.ai's dense GLM-4 decoder (interleaved partial RoPE, joint gate_up MLP)",
100 "Glm4ForCausalLM": "Z.ai's GLM-4-0414 / GLM-Z1 decoder with sandwich norms",
101 "Glm4vForConditionalGeneration": "Z.ai's GLM-4V / GLM-4.1V reasoning vision-language model",
102 "GlmAsrForConditionalGeneration": "Z.ai's GLM-ASR speech recognizer (audio encoder + Llama decoder)",
103 "Starcoder2ForCausalLM": "BigCode's StarCoder2 code model (LayerNorm, biased projections)",
104 "BitNetForCausalLM": "Microsoft's BitNet b1.58 (llama layout with attn/ffn sub-layer norms)",
105 "Exaone4ForCausalLM": "LG AI Research's EXAONE 4.0 (post-norms in residual, per-head QK-norm, hybrid sliding attention)",
106 "BlenderbotForConditionalGeneration": "Meta's Blenderbot dialogue encoder-decoder (asymmetric stacks)",
107 "MiniMaxM2ForCausalLM": "MiniMax's M2 sparse-MoE decoder with sigmoid routing",
108 "ExaoneForCausalLM": "LG AI Research's EXAONE-3.x Korean-English decoder (remote code)",
109 "FalconMambaForCausalLM": "TII's FalconMamba pure-SSM decoder (Mamba-1 with B/C/dt RMS)",
110 "T5Gemma2ForConditionalGeneration": "Google's T5Gemma2 multimodal encoder-decoder model with merged self+cross decoder attention, QK-norm, and dual RoPE (text-only bridge support)",
111 "HrmTextForCausalLM": "Sapient Intelligence's HRM-Text hierarchical two-timescale recurrent model",
112 "BD3LM": "Kuleshov Group's Block Diffusion Language Model (ICLR 2025) for masked text generation",
113 "DreamModel": "HKU-NLP's Dream 7B discrete-diffusion LM (Qwen2.5-based, bidirectional; also Apple DiffuCoder)",
114 "AudioFlamingo3ForConditionalGeneration": "NVIDIA Audio Flamingo 3 audio LALM (Whisper tower + Qwen2.5 LM)",
115 "Emu3ForConditionalGeneration": "BAAI Emu3 unified next-token text+image generation (VQ tokens in a shared vocab)",
116 "FlexOlmoForCausalLM": "AllenAI FlexOlmo federated MoE (merged OLMo-2 experts, inference-time data opt-out)",
117 "MusicFlamingoForConditionalGeneration": "NVIDIA Music Flamingo / Audio Flamingo Next (AF3 + temporal rotary conditioning)",
118 "GiddForDiffusionLM": "GIDD uniform-noise diffusion LM with self-correction (only open non-masked dLM at scale)",
119 "HyenaDNAForCausalLM": "HazyResearch HyenaDNA genomic LM (attention-free Hyena long-conv operator)",
120 "LLaDA2MoeModelLM": "Ant Group LLaDA 2.x masked block-diffusion LM on a DeepSeek-V3-style MoE decoder",
121 "JetMoeForCausalLM": "MIT-IBM JetMoE, the only open at-scale Mixture-of-Attention-heads model",
122 "LagunaForCausalLM": "poolside Laguna coding MoE (heterogeneous per-layer head counts, softplus attention gating)",
123 "ModernBertDecoderForCausalLM": "JHU Ettin decoders (ModernBERT recipe run causally; paired encoder/decoder science suite)",
124 "Jais2ForCausalLM": "G42/Inception Jais 2 Arabic-English family (Nemotron-shaped LN decoder)",
125 "Ministral3ForCausalLM": "Mistral AI Ministral 3 edge models (pure Mistral shape)",
126 "VaultGemmaForCausalLM": "Google VaultGemma, the only fully DP-SGD-pretrained open LLM (Gemma 2 minus post-norms)",
127 "YoutuForCausalLM": "Tencent Youtu-LLM: laptop-scale dense Multi-head Latent Attention (DeepSeek-V2 MLA, all layers dense)",
128 "SwitchTransformersForConditionalGeneration": "Google Switch Transformer, the foundational top-1 capacity-routed MoE (encoder-decoder)",
129 "RwkvForCausalLM": "BlinkDL RWKV-4 WKV linear-attention RNN (Pythia-parallel Pile suite)",
130 "NanoChatForCausalLM": "Karpathy's nanochat education-stack decoder (weightless norms, relu^2 MLP, capped logits)",
131 "HunYuanDenseV1ForCausalLM": "Tencent's open source decoder models",
132 "ViTModel": "Vision Transformer (bare encoder, no classification head)",
133 "ViTForImageClassification": "Vision Transformer with an image classification head",
134 "DeiTModel": "Data-efficient Image Transformer (bare encoder)",
135 "DeiTForImageClassification": "DeiT with a single CLS-token classification head",
136 "Wav2Vec2Model": "Facebook's Wav2Vec 2.0 for speech",
137 "HubertModel": "Facebook's HuBERT for speech",
138 "Cohere2ForCausalLM": "Cohere's Command-A architecture with interleaved sliding-window RoPE and full-attention NoPE layers",
139 "OuroForCausalLM": "ByteDance's Ouro looped language model (LoopLM) with weight-shared iterated depth",
140 "RavenForCausalLM": "tomg-group-umd's Huginn depth-recurrent decoder (prelude / weight-tied recurrent core / coda) with runtime recurrence count",
141 "RWKV7ForCausalLM": 'fla-hub\'s RWKV-7 ("Goose") attention-free recurrent decoder with generalized-delta-rule time-mixing and squared-ReLU channel-mixing',
142 # Unsupported architectures
143 "BertModel": "Google's BERT bidirectional encoder for understanding tasks",
144 "BertForMaskedLM": "BERT with masked language modeling head",
145 "BertForSequenceClassification": "BERT fine-tuned for classification",
146 "RobertaModel": "Facebook's RoBERTa, optimized BERT training",
147 "RobertaForMaskedLM": "RoBERTa with masked language modeling head",
148 "DistilBertModel": "Distilled version of BERT, 40% smaller",
149 "AlbertModel": "A Lite BERT with parameter sharing",
150 "XLNetLMHeadModel": "Google/CMU's XLNet with permutation language modeling",
151 "ElectraModel": "Google's ELECTRA with replaced token detection",
152 "DebertaModel": "Microsoft's DeBERTa with disentangled attention",
153 "DebertaV2Model": "DeBERTa version 2 with improved architecture",
154 "MPNetModel": "Microsoft's MPNet combining MLM and PLM",
155 "LongformerModel": "Allen AI's Longformer for long documents",
156 "BigBirdModel": "Google's BigBird with sparse attention",
157 "ReformerModel": "Google's Reformer with locality-sensitive hashing",
158 "MBartForConditionalGeneration": "Multilingual BART",
159 "PegasusForConditionalGeneration": "Google's PEGASUS for summarization",
160 "MT5ForConditionalGeneration": "Multilingual T5",
161 "WhisperForConditionalGeneration": "OpenAI's Whisper speech recognition",
162 "CLIPModel": "OpenAI's CLIP vision-language model",
163 "SwinModel": "Microsoft's Swin Transformer for vision",
164 "BeitModel": "Microsoft's BERT pre-training for images",
165 "ConvNextModel": "Facebook's ConvNeXt modernized ConvNet",
166 "SegformerModel": "NVIDIA's SegFormer for segmentation",
167 "SpeechT5Model": "Microsoft's SpeechT5 for speech tasks",
168 "BlipModel": "Salesforce's BLIP vision-language model",
169 "Blip2Model": "Salesforce's BLIP-2 with frozen LLM",
170 "LlavaForConditionalGeneration": "Visual instruction-tuned LLaMA",
171 "GitModel": "Microsoft's GIT for vision-language",
172 "PaliGemmaForConditionalGeneration": "Google's PaliGemma vision-language",
173 "CohereForCausalLM": "Cohere's Command models",
174 "DeepseekForCausalLM": "DeepSeek's open models",
175 "InternLMForCausalLM": "Shanghai AI Lab's InternLM",
176 "BaichuanForCausalLM": "Baichuan's Chinese-focused models",
177 "YiForCausalLM": "01.AI's Yi model series",
178 "OrionForCausalLM": "OrionStar's Orion models",
179 "StarcoderForCausalLM": "BigCode's StarCoder for code",
180 "CodeLlamaForCausalLM": "Meta's Code Llama for programming",
181 "CodeGenForCausalLM": "Salesforce's CodeGen models",
182 "SantacoderForCausalLM": "BigCode's SantaCoder",
183}
186def get_architecture_description(arch_id: str) -> str:
187 """Get a description for an architecture, with fallback."""
188 if arch_id in ARCHITECTURE_DESCRIPTIONS:
189 return ARCHITECTURE_DESCRIPTIONS[arch_id]
191 # Generate a basic description from the name
192 if "ForCausalLM" in arch_id:
193 base = arch_id.replace("ForCausalLM", "")
194 return f"{base} architecture for causal language modeling"
195 elif "ForConditionalGeneration" in arch_id:
196 base = arch_id.replace("ForConditionalGeneration", "")
197 return f"{base} encoder-decoder for conditional generation"
198 elif "ForMaskedLM" in arch_id:
199 base = arch_id.replace("ForMaskedLM", "")
200 return f"{base} with masked language modeling head"
201 elif "ForSequenceClassification" in arch_id:
202 base = arch_id.replace("ForSequenceClassification", "")
203 return f"{base} fine-tuned for sequence classification"
204 elif "Model" in arch_id:
205 base = arch_id.replace("Model", "")
206 return f"{base} base model architecture"
207 else:
208 return "Transformer architecture"
211def generate_report(output_path: Path | None = None) -> str:
212 """Generate the markdown report.
214 Args:
215 output_path: Optional path to write the report. If None, only returns the string.
217 Returns:
218 The generated markdown report as a string.
219 """
220 # Gather data
221 models = get_supported_models()
222 architectures = get_supported_architectures()
223 gaps = get_unsupported_architectures()
224 stats = get_registry_stats()
226 # Group models by architecture
227 models_by_arch: dict[str, list[str]] = {}
228 for model in models:
229 arch = model.architecture_id
230 if arch not in models_by_arch:
231 models_by_arch[arch] = []
232 models_by_arch[arch].append(model.model_id)
234 # Sort models within each architecture
235 for arch in models_by_arch:
236 models_by_arch[arch].sort()
238 # Calculate totals
239 total_supported = len(models)
240 total_unsupported = sum(g.total_models for g in gaps)
241 total_all = total_supported + total_unsupported
243 # Build report
244 lines = []
245 lines.append("# TransformerLens Model Compatibility Report")
246 lines.append("")
247 lines.append(f"*Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*")
248 lines.append("")
250 # Summary
251 lines.append("## Summary")
252 lines.append("")
253 lines.append(f"| Metric | Count |")
254 lines.append(f"|--------|-------|")
255 lines.append(f"| Supported Models | {total_supported:,} |")
256 lines.append(f"| Supported Architectures | {len(architectures)} |")
257 lines.append(f"| Verified Models | {stats['total_verified']} |")
258 lines.append(f"| Provisional Models | {stats.get('total_provisional', 0)} |")
259 lines.append(f"| Unsupported Architectures | {len(gaps)} |")
260 lines.append(f"| Models in Unsupported Architectures | {total_unsupported:,} |")
261 lines.append(f"| **Total Potential Models** | **{total_all:,}** |")
262 lines.append("")
264 # Supported models section
265 lines.append("## Supported Models")
266 lines.append("")
267 lines.append(
268 f"TransformerLens supports **{total_supported:,} models** across **{len(architectures)} architectures**."
269 )
270 lines.append("")
272 for arch in sorted(models_by_arch.keys()):
273 model_list = models_by_arch[arch]
274 desc = get_architecture_description(arch)
275 lines.append(f"### {arch}")
276 lines.append("")
277 lines.append(f"*{desc}*")
278 lines.append("")
279 lines.append(f"**{len(model_list)} models:**")
280 lines.append("")
281 for model_id in model_list:
282 model_entry = next((m for m in models if m.model_id == model_id), None)
283 if model_entry and model_entry.status == 1:
284 badge = " ✓"
285 elif model_entry and model_entry.status == 4:
286 badge = " ⚠"
287 else:
288 badge = ""
289 lines.append(f"- `{model_id}`{badge}")
290 lines.append("")
292 # Unsupported architectures section
293 lines.append("## Unsupported Architectures")
294 lines.append("")
295 lines.append(
296 f"The following **{len(gaps)} architectures** are not yet supported by TransformerLens,"
297 )
298 lines.append(f"representing **{total_unsupported:,} models** on HuggingFace.")
299 lines.append("")
300 lines.append("| Architecture | Models | Description |")
301 lines.append("|--------------|--------|-------------|")
303 for gap in gaps:
304 desc = get_architecture_description(gap.architecture_id)
305 lines.append(f"| `{gap.architecture_id}` | {gap.total_models:,} | {desc} |")
307 lines.append("")
309 # Footer
310 lines.append("---")
311 lines.append("")
312 lines.append(
313 "*Report generated by `python -m transformer_lens.tools.model_registry.generate_report`*"
314 )
315 lines.append("")
316 lines.append("✓ = Verified to work with TransformerLens")
317 lines.append(
318 "⚠ = Provisional (structural check only; not numerically verified against HuggingFace)"
319 )
321 report = "\n".join(lines)
323 # Write to file if path provided
324 if output_path:
325 output_path.write_text(report)
326 print(f"Report written to: {output_path}")
328 return report
331def main():
332 """CLI entry point."""
333 parser = argparse.ArgumentParser(
334 description="Generate a markdown report of TransformerLens model compatibility.",
335 formatter_class=argparse.RawDescriptionHelpFormatter,
336 epilog="""
337Examples:
338 # Generate report to default location (MODEL_COMPATIBILITY_REPORT.md)
339 python -m transformer_lens.tools.model_registry.generate_report
341 # Generate report to custom location
342 python -m transformer_lens.tools.model_registry.generate_report -o my_report.md
344 # Print report to stdout only
345 python -m transformer_lens.tools.model_registry.generate_report --stdout
346""",
347 )
348 parser.add_argument(
349 "-o",
350 "--output",
351 type=Path,
352 default=None,
353 help="Output file path (default: MODEL_COMPATIBILITY_REPORT.md in current directory)",
354 )
355 parser.add_argument(
356 "--stdout",
357 action="store_true",
358 help="Print report to stdout instead of writing to file",
359 )
361 args = parser.parse_args()
363 if args.stdout:
364 report = generate_report()
365 print(report)
366 else:
367 output_path = args.output or Path("MODEL_COMPATIBILITY_REPORT.md")
368 generate_report(output_path)
371if __name__ == "__main__": 371 ↛ 372line 371 didn't jump to line 372 because the condition on line 371 was never true
372 main()