Coverage for transformer_lens/utilities/architectures.py: 89%
58 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Centralized architecture classification for TransformerLens.
3Single source of truth for architecture type detection. Used by the bridge
4loading pipeline, benchmarks, and verification tools.
5"""
7from typing import Optional
9# Encoder-decoder models (T5, BART, etc.)
10SEQ2SEQ_ARCHITECTURES: set[str] = {
11 "T5ForConditionalGeneration",
12 "MT5ForConditionalGeneration",
13 "T5WithLMHeadModel",
14 "T5GemmaForConditionalGeneration",
15 "LongT5ForConditionalGeneration",
16 "T5Gemma2ForConditionalGeneration",
17 "BartForConditionalGeneration",
18 "MBartForConditionalGeneration",
19 "M2M100ForConditionalGeneration",
20 "MarianMTModel",
21 "PegasusForConditionalGeneration",
22 "BlenderbotForConditionalGeneration",
23 "BlenderbotSmallForConditionalGeneration",
24 "LEDForConditionalGeneration",
25 "SwitchTransformersForConditionalGeneration",
26}
28# Post-norm decoders: ln1/ln2 normalize each sublayer's OUTPUT before the residual
29# add, so LN folding and writing-weight centering (which assume the gain sits on a
30# sublayer's INPUT) are not valid algebra for them.
31POST_NORM_ARCHITECTURES: set[str] = {
32 "Olmo2ForCausalLM",
33 "Olmo3ForCausalLM",
34}
36# Masked language models (BERT-style, no text generation)
37MASKED_LM_ARCHITECTURES: set[str] = {
38 "BertForMaskedLM",
39 "RobertaForMaskedLM",
40 "AlbertForMaskedLM",
41 "DistilBertForMaskedLM",
42 "ElectraForMaskedLM",
43 "BD3LM",
44}
46# Vision-language multimodal models
47MULTIMODAL_ARCHITECTURES: set[str] = {
48 "Emu3ForConditionalGeneration",
49 "LlavaForConditionalGeneration",
50 "LlavaNextForConditionalGeneration",
51 "LlavaOnevisionForConditionalGeneration",
52 "Gemma3ForConditionalGeneration",
53 "Gemma4ForConditionalGeneration",
54 "Qwen3_5ForConditionalGeneration",
55 "Qwen3_5MoeForConditionalGeneration",
56 "Idefics3ForConditionalGeneration",
57 "Florence2ForConditionalGeneration",
58 "Mistral3ForConditionalGeneration",
59 "Llama4ForConditionalGeneration",
60 "Qwen2_5_VLForConditionalGeneration",
61 "Qwen3VLForConditionalGeneration",
62 "Qwen3VLMoeForConditionalGeneration",
63 "Glm4vForConditionalGeneration",
64}
66# Audio-conditioned text decoders (audio encoder + causal LM); load via
67# AutoModelForSeq2SeqLM but behave as text decoders for classification.
68AUDIO_TEXT_ARCHITECTURES: set[str] = {
69 "Qwen2AudioForConditionalGeneration",
70 "GlmAsrForConditionalGeneration",
71 "AudioFlamingo3ForConditionalGeneration",
72 "MusicFlamingoForConditionalGeneration",
73}
75# Audio spectrogram models for classification
76AUDIO_CLASSIFICATION_ARCHITECTURES: set[str] = {
77 "ASTForAudioClassification",
78}
80# Audio encoder models (HuBERT, wav2vec2, etc.)
81AUDIO_ARCHITECTURES: set[str] = {
82 "HubertForCTC",
83 "HubertModel",
84 "HubertForSequenceClassification",
85} | AUDIO_CLASSIFICATION_ARCHITECTURES
87# Vision-only (non-multimodal, no text tower) encoder models. Split into the
88# two HF AutoModel classes they load under: bare encoders load via AutoModel,
89# classification heads load via AutoModelForImageClassification.
90VISION_MODEL_ARCHITECTURES: set[str] = {
91 "ViTModel",
92 "DeiTModel",
93}
94VISION_CLASSIFICATION_ARCHITECTURES: set[str] = {
95 "ViTForImageClassification",
96 "DeiTForImageClassification",
97}
98VISION_ARCHITECTURES: set[str] = VISION_MODEL_ARCHITECTURES | VISION_CLASSIFICATION_ARCHITECTURES
100# Text models whose remote code registers only under plain AutoModel
101# (the class itself carries the LM head).
102BASE_AUTOMODEL_ARCHITECTURES: set[str] = {
103 "DreamModel",
104}
106# Bridge uses different hook shapes than HookedTransformer by design.
107# Phase 2/3 HT comparisons are skipped; Phase 1 (HF comparison) is the gold standard.
108NO_HT_COMPARISON_ARCHITECTURES: set[str] = (
109 MULTIMODAL_ARCHITECTURES
110 | AUDIO_ARCHITECTURES
111 # Vision encoders have no HookedTransformer counterpart.
112 | VISION_ARCHITECTURES
113 # Encoder-decoder: HookedTransformer cannot represent them (T5 repos under
114 # org-prefixed names slip past HT's legacy name guard and crash at forward).
115 | SEQ2SEQ_ARCHITECTURES
116 | {
117 "Gemma3ForCausalLM",
118 }
119)
122def classify_architecture(architecture: str) -> str:
123 """Classify an architecture string into a model type.
125 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm"
126 """
127 if architecture in SEQ2SEQ_ARCHITECTURES:
128 return "seq2seq"
129 if architecture in MASKED_LM_ARCHITECTURES:
130 return "masked_lm"
131 if architecture in MULTIMODAL_ARCHITECTURES:
132 return "multimodal"
133 if architecture in AUDIO_ARCHITECTURES:
134 return "audio"
135 if architecture in VISION_ARCHITECTURES:
136 return "vision"
137 return "causal_lm"
140def get_architectures_for_config(config) -> list[str]:
141 """Extract architecture strings from an HF config object."""
142 architectures = []
143 if hasattr(config, "original_architecture"): 143 ↛ 144line 143 didn't jump to line 144 because the condition on line 143 was never true
144 architectures.append(config.original_architecture)
145 if hasattr(config, "architectures") and config.architectures: 145 ↛ 147line 145 didn't jump to line 147 because the condition on line 145 was always true
146 architectures.extend(config.architectures)
147 return architectures
150def classify_model_config(config) -> str:
151 """Classify a model by its HF config.
153 Checks config.is_encoder_decoder first, then falls back to architecture list.
154 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "causal_lm"
155 """
156 if getattr(config, "is_encoder_decoder", False): 156 ↛ 157line 156 didn't jump to line 157 because the condition on line 156 was never true
157 return "seq2seq"
158 for arch in get_architectures_for_config(config): 158 ↛ 162line 158 didn't jump to line 162 because the loop on line 158 didn't complete
159 model_type = classify_architecture(arch)
160 if model_type != "causal_lm": 160 ↛ 158line 160 didn't jump to line 158 because the condition on line 160 was always true
161 return model_type
162 return "causal_lm"
165def classify_model_name(
166 model_name: str,
167 trust_remote_code: bool = False,
168 token: Optional[str] = None,
169) -> str:
170 """Classify a model by its HuggingFace model name.
172 Loads the config once, classifies from it. If token is None, reads
173 HF_TOKEN from the environment automatically.
174 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "causal_lm"
175 """
176 try:
177 from transformers import AutoConfig
179 if token is None: 179 ↛ 184line 179 didn't jump to line 184 because the condition on line 179 was always true
180 from transformer_lens.utilities.hf_utils import get_hf_token
182 token = get_hf_token()
184 config = AutoConfig.from_pretrained(
185 model_name, trust_remote_code=trust_remote_code, token=token
186 )
187 return classify_model_config(config)
188 except Exception:
189 return "causal_lm"
192def is_masked_lm_model(
193 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
194) -> bool:
195 """Check if a model is a masked language model (BERT-style)."""
196 return (
197 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
198 == "masked_lm"
199 )
202def is_encoder_decoder_model(
203 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
204) -> bool:
205 """Check if a model is an encoder-decoder architecture (T5, BART, etc.)."""
206 return (
207 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
208 == "seq2seq"
209 )
212def is_multimodal_model(
213 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
214) -> bool:
215 """Check if a model is a multimodal vision-language model (LLaVA, Gemma3)."""
216 return (
217 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
218 == "multimodal"
219 )
222def is_audio_model(
223 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
224) -> bool:
225 """Check if a model is an audio encoder model (HuBERT, wav2vec2)."""
226 return (
227 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token) == "audio"
228 )