Coverage for transformer_lens/tools/model_registry/generate_report.py: 0%
107 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-06-09 00:32 +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 "GPT2LMHeadModel": "OpenAI's GPT-2 decoder-only transformer for causal language modeling",
30 "GPTNeoForCausalLM": "EleutherAI's GPT-Neo, an open-source GPT-3-like model",
31 "GPTNeoXForCausalLM": "EleutherAI's GPT-NeoX architecture used in Pythia models",
32 "GPTJForCausalLM": "EleutherAI's GPT-J 6B parameter model",
33 "LlamaForCausalLM": "Meta's LLaMA architecture, basis for many open models",
34 "MistralForCausalLM": "Mistral AI's efficient 7B parameter model with sliding window attention",
35 "MixtralForCausalLM": "Mistral AI's Mixture of Experts model",
36 "GemmaForCausalLM": "Google's Gemma lightweight open model family",
37 "Gemma2ForCausalLM": "Google's Gemma 2 with improved architecture",
38 "Gemma3ForCausalLM": "Google's Gemma 3 latest generation",
39 "Gemma3nForConditionalGeneration": "Google's Gemma 3n efficient tri-modal model (text-only support)",
40 "Qwen2ForCausalLM": "Alibaba's Qwen2 multilingual model",
41 "Qwen3ForCausalLM": "Alibaba's Qwen3 latest generation",
42 "Qwen3_5ForConditionalGeneration": "Alibaba's Qwen3.5 vision-language model",
43 "BloomForCausalLM": "BigScience's BLOOM multilingual model",
44 "OPTForCausalLM": "Meta's Open Pre-trained Transformer",
45 "PhiForCausalLM": "Microsoft's Phi small language model",
46 "Phi3ForCausalLM": "Microsoft's Phi-3 improved small model",
47 "FalconForCausalLM": "TII's Falcon model series",
48 "OlmoForCausalLM": "Allen AI's OLMo open language model",
49 "Olmo2ForCausalLM": "Allen AI's OLMo 2 with improved training",
50 "Olmo3ForCausalLM": "Allen AI's OLMo 3 latest generation",
51 "OlmoeForCausalLM": "Allen AI's OLMoE Mixture of Experts model",
52 "StableLmForCausalLM": "Stability AI's StableLM model",
53 "SmolLM3ForCausalLM": "Hugging Face's SmolLM3 compact open model with NoPE layers",
54 "T5ForConditionalGeneration": "Google's T5 encoder-decoder model (partial support)",
55 # Unsupported architectures
56 "BertModel": "Google's BERT bidirectional encoder for understanding tasks",
57 "BertForMaskedLM": "BERT with masked language modeling head",
58 "BertForSequenceClassification": "BERT fine-tuned for classification",
59 "RobertaModel": "Facebook's RoBERTa, optimized BERT training",
60 "RobertaForMaskedLM": "RoBERTa with masked language modeling head",
61 "DistilBertModel": "Distilled version of BERT, 40% smaller",
62 "AlbertModel": "A Lite BERT with parameter sharing",
63 "XLNetLMHeadModel": "Google/CMU's XLNet with permutation language modeling",
64 "ElectraModel": "Google's ELECTRA with replaced token detection",
65 "DebertaModel": "Microsoft's DeBERTa with disentangled attention",
66 "DebertaV2Model": "DeBERTa version 2 with improved architecture",
67 "MPNetModel": "Microsoft's MPNet combining MLM and PLM",
68 "LongformerModel": "Allen AI's Longformer for long documents",
69 "BigBirdModel": "Google's BigBird with sparse attention",
70 "ReformerModel": "Google's Reformer with locality-sensitive hashing",
71 "BartForConditionalGeneration": "Facebook's BART encoder-decoder model",
72 "MBartForConditionalGeneration": "Multilingual BART",
73 "PegasusForConditionalGeneration": "Google's PEGASUS for summarization",
74 "MT5ForConditionalGeneration": "Multilingual T5",
75 "WhisperForConditionalGeneration": "OpenAI's Whisper speech recognition",
76 "CLIPModel": "OpenAI's CLIP vision-language model",
77 "ViTModel": "Google's Vision Transformer",
78 "SwinModel": "Microsoft's Swin Transformer for vision",
79 "DeiTModel": "Facebook's Data-efficient Image Transformer",
80 "BeitModel": "Microsoft's BERT pre-training for images",
81 "ConvNextModel": "Facebook's ConvNeXt modernized ConvNet",
82 "SegformerModel": "NVIDIA's SegFormer for segmentation",
83 "Wav2Vec2Model": "Facebook's Wav2Vec 2.0 for speech",
84 "HubertModel": "Facebook's HuBERT for speech",
85 "SpeechT5Model": "Microsoft's SpeechT5 for speech tasks",
86 "BlipModel": "Salesforce's BLIP vision-language model",
87 "Blip2Model": "Salesforce's BLIP-2 with frozen LLM",
88 "LlavaForConditionalGeneration": "Visual instruction-tuned LLaMA",
89 "GitModel": "Microsoft's GIT for vision-language",
90 "PaliGemmaForConditionalGeneration": "Google's PaliGemma vision-language",
91 "CohereForCausalLM": "Cohere's Command models",
92 "DeepseekForCausalLM": "DeepSeek's open models",
93 "InternLMForCausalLM": "Shanghai AI Lab's InternLM",
94 "BaichuanForCausalLM": "Baichuan's Chinese-focused models",
95 "YiForCausalLM": "01.AI's Yi model series",
96 "OrionForCausalLM": "OrionStar's Orion models",
97 "StarcoderForCausalLM": "BigCode's StarCoder for code",
98 "CodeLlamaForCausalLM": "Meta's Code Llama for programming",
99 "CodeGenForCausalLM": "Salesforce's CodeGen models",
100 "SantacoderForCausalLM": "BigCode's SantaCoder",
101}
104def get_architecture_description(arch_id: str) -> str:
105 """Get a description for an architecture, with fallback."""
106 if arch_id in ARCHITECTURE_DESCRIPTIONS:
107 return ARCHITECTURE_DESCRIPTIONS[arch_id]
109 # Generate a basic description from the name
110 if "ForCausalLM" in arch_id:
111 base = arch_id.replace("ForCausalLM", "")
112 return f"{base} architecture for causal language modeling"
113 elif "ForConditionalGeneration" in arch_id:
114 base = arch_id.replace("ForConditionalGeneration", "")
115 return f"{base} encoder-decoder for conditional generation"
116 elif "ForMaskedLM" in arch_id:
117 base = arch_id.replace("ForMaskedLM", "")
118 return f"{base} with masked language modeling head"
119 elif "ForSequenceClassification" in arch_id:
120 base = arch_id.replace("ForSequenceClassification", "")
121 return f"{base} fine-tuned for sequence classification"
122 elif "Model" in arch_id:
123 base = arch_id.replace("Model", "")
124 return f"{base} base model architecture"
125 else:
126 return "Transformer architecture"
129def generate_report(output_path: Path | None = None) -> str:
130 """Generate the markdown report.
132 Args:
133 output_path: Optional path to write the report. If None, only returns the string.
135 Returns:
136 The generated markdown report as a string.
137 """
138 # Gather data
139 models = get_supported_models()
140 architectures = get_supported_architectures()
141 gaps = get_unsupported_architectures()
142 stats = get_registry_stats()
144 # Group models by architecture
145 models_by_arch: dict[str, list[str]] = {}
146 for model in models:
147 arch = model.architecture_id
148 if arch not in models_by_arch:
149 models_by_arch[arch] = []
150 models_by_arch[arch].append(model.model_id)
152 # Sort models within each architecture
153 for arch in models_by_arch:
154 models_by_arch[arch].sort()
156 # Calculate totals
157 total_supported = len(models)
158 total_unsupported = sum(g.total_models for g in gaps)
159 total_all = total_supported + total_unsupported
161 # Build report
162 lines = []
163 lines.append("# TransformerLens Model Compatibility Report")
164 lines.append("")
165 lines.append(f"*Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*")
166 lines.append("")
168 # Summary
169 lines.append("## Summary")
170 lines.append("")
171 lines.append(f"| Metric | Count |")
172 lines.append(f"|--------|-------|")
173 lines.append(f"| Supported Models | {total_supported:,} |")
174 lines.append(f"| Supported Architectures | {len(architectures)} |")
175 lines.append(f"| Verified Models | {stats['total_verified']} |")
176 lines.append(f"| Unsupported Architectures | {len(gaps)} |")
177 lines.append(f"| Models in Unsupported Architectures | {total_unsupported:,} |")
178 lines.append(f"| **Total Potential Models** | **{total_all:,}** |")
179 lines.append("")
181 # Supported models section
182 lines.append("## Supported Models")
183 lines.append("")
184 lines.append(
185 f"TransformerLens supports **{total_supported:,} models** across **{len(architectures)} architectures**."
186 )
187 lines.append("")
189 for arch in sorted(models_by_arch.keys()):
190 model_list = models_by_arch[arch]
191 desc = get_architecture_description(arch)
192 lines.append(f"### {arch}")
193 lines.append("")
194 lines.append(f"*{desc}*")
195 lines.append("")
196 lines.append(f"**{len(model_list)} models:**")
197 lines.append("")
198 for model_id in model_list:
199 # Check if verified
200 model_entry = next((m for m in models if m.model_id == model_id), None)
201 verified_badge = " ✓" if model_entry and model_entry.status == 1 else ""
202 lines.append(f"- `{model_id}`{verified_badge}")
203 lines.append("")
205 # Unsupported architectures section
206 lines.append("## Unsupported Architectures")
207 lines.append("")
208 lines.append(
209 f"The following **{len(gaps)} architectures** are not yet supported by TransformerLens,"
210 )
211 lines.append(f"representing **{total_unsupported:,} models** on HuggingFace.")
212 lines.append("")
213 lines.append("| Architecture | Models | Description |")
214 lines.append("|--------------|--------|-------------|")
216 for gap in gaps:
217 desc = get_architecture_description(gap.architecture_id)
218 lines.append(f"| `{gap.architecture_id}` | {gap.total_models:,} | {desc} |")
220 lines.append("")
222 # Footer
223 lines.append("---")
224 lines.append("")
225 lines.append(
226 "*Report generated by `python -m transformer_lens.tools.model_registry.generate_report`*"
227 )
228 lines.append("")
229 lines.append("✓ = Verified to work with TransformerLens")
231 report = "\n".join(lines)
233 # Write to file if path provided
234 if output_path:
235 output_path.write_text(report)
236 print(f"Report written to: {output_path}")
238 return report
241def main():
242 """CLI entry point."""
243 parser = argparse.ArgumentParser(
244 description="Generate a markdown report of TransformerLens model compatibility.",
245 formatter_class=argparse.RawDescriptionHelpFormatter,
246 epilog="""
247Examples:
248 # Generate report to default location (MODEL_COMPATIBILITY_REPORT.md)
249 python -m transformer_lens.tools.model_registry.generate_report
251 # Generate report to custom location
252 python -m transformer_lens.tools.model_registry.generate_report -o my_report.md
254 # Print report to stdout only
255 python -m transformer_lens.tools.model_registry.generate_report --stdout
256""",
257 )
258 parser.add_argument(
259 "-o",
260 "--output",
261 type=Path,
262 default=None,
263 help="Output file path (default: MODEL_COMPATIBILITY_REPORT.md in current directory)",
264 )
265 parser.add_argument(
266 "--stdout",
267 action="store_true",
268 help="Print report to stdout instead of writing to file",
269 )
271 args = parser.parse_args()
273 if args.stdout:
274 report = generate_report()
275 print(report)
276 else:
277 output_path = args.output or Path("MODEL_COMPATIBILITY_REPORT.md")
278 generate_report(output_path)
281if __name__ == "__main__":
282 main()