Coverage for transformer_lens/utilities/architectures.py: 82%
57 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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 "Wav2Vec2ForCTC",
86 "Wav2Vec2Model",
87 # Pretraining checkpoints (facebook/wav2vec2-base/-large declare this class)
88 # load their encoder via AutoModel -> Wav2Vec2Model.
89 "Wav2Vec2ForPreTraining",
90} | AUDIO_CLASSIFICATION_ARCHITECTURES
92# Vision-only (non-multimodal, no text tower) encoder models. Split into the
93# two HF AutoModel classes they load under: bare encoders load via AutoModel,
94# classification heads load via AutoModelForImageClassification.
95VISION_MODEL_ARCHITECTURES: set[str] = {
96 "ViTModel",
97 "DeiTModel",
98}
99VISION_CLASSIFICATION_ARCHITECTURES: set[str] = {
100 "ViTForImageClassification",
101 "DeiTForImageClassification",
102}
103VISION_ARCHITECTURES: set[str] = VISION_MODEL_ARCHITECTURES | VISION_CLASSIFICATION_ARCHITECTURES
105# Text models whose remote code registers only under plain AutoModel
106# (the class itself carries the LM head).
107BASE_AUTOMODEL_ARCHITECTURES: set[str] = {
108 "DreamModel",
109}
112def classify_architecture(architecture: str) -> str:
113 """Classify an architecture string into a model type.
115 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm"
116 """
117 if architecture in SEQ2SEQ_ARCHITECTURES:
118 return "seq2seq"
119 if architecture in MASKED_LM_ARCHITECTURES:
120 return "masked_lm"
121 if architecture in MULTIMODAL_ARCHITECTURES:
122 return "multimodal"
123 if architecture in AUDIO_ARCHITECTURES:
124 return "audio"
125 if architecture in VISION_ARCHITECTURES:
126 return "vision"
127 return "causal_lm"
130def get_architectures_for_config(config) -> list[str]:
131 """Extract architecture strings from an HF config object."""
132 architectures = []
133 if hasattr(config, "original_architecture"): 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true
134 architectures.append(config.original_architecture)
135 if hasattr(config, "architectures") and config.architectures: 135 ↛ 137line 135 didn't jump to line 137 because the condition on line 135 was always true
136 architectures.extend(config.architectures)
137 return architectures
140def classify_model_config(config) -> str:
141 """Classify a model by its HF config.
143 Checks config.is_encoder_decoder first, then falls back to architecture list.
144 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm"
145 """
146 if getattr(config, "is_encoder_decoder", False): 146 ↛ 147line 146 didn't jump to line 147 because the condition on line 146 was never true
147 return "seq2seq"
148 for arch in get_architectures_for_config(config): 148 ↛ 152line 148 didn't jump to line 152 because the loop on line 148 didn't complete
149 model_type = classify_architecture(arch)
150 if model_type != "causal_lm": 150 ↛ 148line 150 didn't jump to line 148 because the condition on line 150 was always true
151 return model_type
152 return "causal_lm"
155def classify_model_name(
156 model_name: str,
157 trust_remote_code: bool = False,
158 token: Optional[str] = None,
159) -> str:
160 """Classify a model by its HuggingFace model name.
162 Loads the config once, classifies from it. If token is None, reads
163 HF_TOKEN from the environment automatically.
164 Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm"
165 """
166 try:
167 from transformers import AutoConfig
169 if token is None: 169 ↛ 174line 169 didn't jump to line 174 because the condition on line 169 was always true
170 from transformer_lens.utilities.hf_utils import get_hf_token
172 token = get_hf_token()
174 config = AutoConfig.from_pretrained(
175 model_name, trust_remote_code=trust_remote_code, token=token
176 )
177 return classify_model_config(config)
178 except Exception:
179 return "causal_lm"
182def is_masked_lm_model(
183 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
184) -> bool:
185 """Check if a model is a masked language model (BERT-style)."""
186 return (
187 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
188 == "masked_lm"
189 )
192def is_encoder_decoder_model(
193 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
194) -> bool:
195 """Check if a model is an encoder-decoder architecture (T5, BART, etc.)."""
196 return (
197 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
198 == "seq2seq"
199 )
202def is_multimodal_model(
203 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
204) -> bool:
205 """Check if a model is a multimodal vision-language model (LLaVA, Gemma3)."""
206 return (
207 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token)
208 == "multimodal"
209 )
212def is_audio_model(
213 model_name: str, trust_remote_code: bool = False, token: Optional[str] = None
214) -> bool:
215 """Check if a model is an audio encoder model (HuBERT, wav2vec2)."""
216 return (
217 classify_model_name(model_name, trust_remote_code=trust_remote_code, token=token) == "audio"
218 )