Coverage for transformer_lens/utilities/architectures.py: 82%
57 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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 "MuseGlimmerForConditionalGeneration",
61 "Qwen2_5_VLForConditionalGeneration",
62 "Qwen3VLForConditionalGeneration",
63 "Qwen3VLMoeForConditionalGeneration",
64 "Glm4vForConditionalGeneration",
65}
67# Audio-conditioned text decoders (audio encoder + causal LM); load via
68# AutoModelForSeq2SeqLM but behave as text decoders for classification.
69AUDIO_TEXT_ARCHITECTURES: set[str] = {
70 "Qwen2AudioForConditionalGeneration",
71 "GlmAsrForConditionalGeneration",
72 "AudioFlamingo3ForConditionalGeneration",
73 "MusicFlamingoForConditionalGeneration",
74}
76# Audio spectrogram models for classification
77AUDIO_CLASSIFICATION_ARCHITECTURES: set[str] = {
78 "ASTForAudioClassification",
79}
81# Audio encoder models (HuBERT, wav2vec2, etc.)
82AUDIO_ARCHITECTURES: set[str] = {
83 "HubertForCTC",
84 "HubertModel",
85 "HubertForSequenceClassification",
86 "Wav2Vec2ForCTC",
87 "Wav2Vec2Model",
88 # Pretraining checkpoints (facebook/wav2vec2-base/-large declare this class)
89 # load their encoder via AutoModel -> Wav2Vec2Model.
90 "Wav2Vec2ForPreTraining",
91} | AUDIO_CLASSIFICATION_ARCHITECTURES
93# Vision-only (non-multimodal, no text tower) encoder models. Split into the
94# two HF AutoModel classes they load under: bare encoders load via AutoModel,
95# classification heads load via AutoModelForImageClassification.
96VISION_MODEL_ARCHITECTURES: set[str] = {
97 "ViTModel",
98 "DeiTModel",
99}
100VISION_CLASSIFICATION_ARCHITECTURES: set[str] = {
101 "ViTForImageClassification",
102 "DeiTForImageClassification",
103}
104VISION_ARCHITECTURES: set[str] = VISION_MODEL_ARCHITECTURES | VISION_CLASSIFICATION_ARCHITECTURES
106# Text models whose remote code registers only under plain AutoModel
107# (the class itself carries the LM head).
108BASE_AUTOMODEL_ARCHITECTURES: set[str] = {
109 "DreamModel",
110}
113def classify_architecture(architecture: str) -> str:
114 """Classify an architecture string into a model type.
116 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm"
117 """
118 if architecture in SEQ2SEQ_ARCHITECTURES:
119 return "seq2seq"
120 if architecture in MASKED_LM_ARCHITECTURES:
121 return "masked_lm"
122 if architecture in MULTIMODAL_ARCHITECTURES:
123 return "multimodal"
124 if architecture in AUDIO_ARCHITECTURES:
125 return "audio"
126 if architecture in VISION_ARCHITECTURES:
127 return "vision"
128 return "causal_lm"
131def get_architectures_for_config(config) -> list[str]:
132 """Extract architecture strings from an HF config object."""
133 architectures = []
134 if hasattr(config, "original_architecture"): 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true
135 architectures.append(config.original_architecture)
136 if hasattr(config, "architectures") and config.architectures: 136 ↛ 138line 136 didn't jump to line 138 because the condition on line 136 was always true
137 architectures.extend(config.architectures)
138 return architectures
141def classify_model_config(config) -> str:
142 """Classify a model by its HF config.
144 Checks config.is_encoder_decoder first, then falls back to architecture list.
145 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm"
146 """
147 if getattr(config, "is_encoder_decoder", False): 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true
148 return "seq2seq"
149 for arch in get_architectures_for_config(config): 149 ↛ 153line 149 didn't jump to line 153 because the loop on line 149 didn't complete
150 model_type = classify_architecture(arch)
151 if model_type != "causal_lm": 151 ↛ 149line 151 didn't jump to line 149 because the condition on line 151 was always true
152 return model_type
153 return "causal_lm"
156def classify_model_name(
157 model_name: str,
158 trust_remote_code: bool = False,
159 token: Optional[str] = None,
160) -> str:
161 """Classify a model by its HuggingFace model name.
163 Loads the config once, classifies from it. If token is None, reads
164 HF_TOKEN from the environment automatically.
165 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm"
166 """
167 try:
168 from transformers import AutoConfig
170 if token is None: 170 ↛ 175line 170 didn't jump to line 175 because the condition on line 170 was always true
171 from transformer_lens.utilities.hf_utils import get_hf_token
173 token = get_hf_token()
175 config = AutoConfig.from_pretrained(
176 model_name, trust_remote_code=trust_remote_code, token=token
177 )
178 return classify_model_config(config)
179 except Exception:
180 return "causal_lm"
183def is_masked_lm_model(
184 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
185) -> bool:
186 """Check if a model is a masked language model (BERT-style)."""
187 return (
188 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
189 == "masked_lm"
190 )
193def is_encoder_decoder_model(
194 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
195) -> bool:
196 """Check if a model is an encoder-decoder architecture (T5, BART, etc.)."""
197 return (
198 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
199 == "seq2seq"
200 )
203def is_multimodal_model(
204 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
205) -> bool:
206 """Check if a model is a multimodal vision-language model (LLaVA, Gemma3)."""
207 return (
208 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
209 == "multimodal"
210 )
213def is_audio_model(
214 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
215) -> bool:
216 """Check if a model is an audio encoder model (HuBERT, wav2vec2)."""
217 return (
218 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token) == "audio"
219 )