Coverage for transformer_lens/model_bridge/sources/transformers/helpers.py: 67%
108 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""HF-loader-specific helpers: model-class selection, modality processor loading, checkpoint revision resolution, registry discovery."""
2from __future__ import annotations
4from typing import Any
6import transformers
7from transformers import (
8 AutoModelForCausalLM,
9 AutoModelForMaskedLM,
10 AutoModelForSeq2SeqLM,
11)
13from transformer_lens.tools.model_registry.checkpoints import get_checkpoint_labels
16def get_hf_model_class_for_architecture(architecture: str):
17 """Pick the correct HuggingFace ``AutoModel*`` class for the architecture."""
18 from transformer_lens.utilities.architectures import (
19 AUDIO_ARCHITECTURES,
20 AUDIO_CLASSIFICATION_ARCHITECTURES,
21 AUDIO_TEXT_ARCHITECTURES,
22 BASE_AUTOMODEL_ARCHITECTURES,
23 MASKED_LM_ARCHITECTURES,
24 MULTIMODAL_ARCHITECTURES,
25 SEQ2SEQ_ARCHITECTURES,
26 VISION_ARCHITECTURES,
27 VISION_CLASSIFICATION_ARCHITECTURES,
28 )
30 if architecture in SEQ2SEQ_ARCHITECTURES or architecture in AUDIO_TEXT_ARCHITECTURES:
31 return AutoModelForSeq2SeqLM
32 elif architecture in MASKED_LM_ARCHITECTURES:
33 return AutoModelForMaskedLM
34 elif architecture in MULTIMODAL_ARCHITECTURES:
35 from transformers import AutoModelForImageTextToText
37 return AutoModelForImageTextToText
38 elif architecture in BASE_AUTOMODEL_ARCHITECTURES: 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true
39 from transformers import AutoModel
41 return AutoModel
42 elif architecture in AUDIO_CLASSIFICATION_ARCHITECTURES:
43 from transformers import AutoModelForAudioClassification
45 return AutoModelForAudioClassification
46 elif architecture in AUDIO_ARCHITECTURES:
47 if "ForCTC" in architecture: 47 ↛ 48line 47 didn't jump to line 48 because the condition on line 47 was never true
48 from transformers import AutoModelForCTC
50 return AutoModelForCTC
51 from transformers import AutoModel
53 return AutoModel
54 elif architecture in VISION_ARCHITECTURES:
55 if architecture in VISION_CLASSIFICATION_ARCHITECTURES:
56 from transformers import AutoModelForImageClassification
58 return AutoModelForImageClassification
59 from transformers import AutoModel
61 return AutoModel
62 else:
63 return AutoModelForCausalLM
66# Modality flag → HF auto-loader class, applied in order. Last write wins on
67# bridge.processor: multimodal → audio → vision. The flags are disjoint today
68# (vit.py is the only is_visual_model adapter and it doesn't set is_multimodal),
69# but the order must be preserved for future dual-flag adapters.
70_MODALITY_PROCESSOR_LOADERS: list[tuple[str, str]] = [
71 ("is_multimodal", "AutoProcessor"),
72 ("is_audio_model", "AutoFeatureExtractor"),
73 ("is_visual_model", "AutoImageProcessor"),
74]
77def _ensure_torchvision() -> bool:
78 """Import torchvision, installing it on the fly if missing; True when importable."""
79 try:
80 import torchvision # noqa: F401
82 return True
83 except Exception:
84 pass
85 import importlib
86 import shutil
87 import subprocess
88 import sys
90 try:
91 if shutil.which("uv"):
92 subprocess.check_call(["uv", "pip", "install", "torchvision", "-q"])
93 else:
94 subprocess.check_call([sys.executable, "-m", "pip", "install", "torchvision", "-q"])
95 importlib.invalidate_caches()
96 return True
97 except Exception:
98 return False
101def load_modality_processor(
102 bridge: Any,
103 cfg: Any,
104 model_name: str,
105 trust_remote_code: bool,
106 token: str | None,
107) -> None:
108 """Attach the modality preprocessor to ``bridge.processor`` per the cfg's modality flags.
110 Best-effort: each loader failure is swallowed so a missing/broken processor
111 never blocks the boot itself.
112 """
113 for cfg_flag, loader_name in _MODALITY_PROCESSOR_LOADERS:
114 if not getattr(cfg, cfg_flag, False):
115 continue
116 try:
117 loader = getattr(transformers, loader_name)
118 bridge.processor = loader.from_pretrained(
119 model_name,
120 token=token,
121 trust_remote_code=trust_remote_code,
122 )
123 except Exception:
124 # Some AutoProcessors need torchvision (e.g. LlavaOnevision); install and retry.
125 if loader_name != "AutoProcessor" or not _ensure_torchvision():
126 continue
127 try:
128 bridge.processor = transformers.AutoProcessor.from_pretrained(
129 model_name,
130 token=token,
131 trust_remote_code=trust_remote_code,
132 )
133 except Exception:
134 pass
137# Known training-checkpoint revision conventions on HF Hub.
138_CHECKPOINT_REVISION_FORMATS: dict[str, str] = {
139 "EleutherAI/pythia": "step{value}",
140 "stanford-crfm": "checkpoint-{value}",
141}
144def _resolve_checkpoint_to_revision(
145 model_name: str,
146 checkpoint_index: int | None,
147 checkpoint_value: int | None,
148) -> str:
149 """Convert a checkpoint index/value into an HF revision string, validated against ``get_checkpoint_labels``."""
150 if checkpoint_index is None and checkpoint_value is None:
151 raise ValueError("Must specify either checkpoint_index or checkpoint_value.")
153 format_str: str | None = None
154 for prefix, fmt in _CHECKPOINT_REVISION_FORMATS.items():
155 if model_name.startswith(prefix):
156 format_str = fmt
157 break
158 if format_str is None:
159 raise ValueError(
160 f"Model {model_name!r} does not have a known checkpoint revision convention. "
161 f"Pass revision= directly if your model uses HF revisions. Known checkpoint "
162 f"families: {list(_CHECKPOINT_REVISION_FORMATS.keys())}."
163 )
165 labels, _ = get_checkpoint_labels(model_name)
166 if checkpoint_value is not None:
167 if checkpoint_value not in labels:
168 raise ValueError(
169 f"checkpoint_value={checkpoint_value} not in available checkpoints for "
170 f"{model_name!r}. {len(labels)} labels available, "
171 f"first/last: {labels[0]}..{labels[-1]}."
172 )
173 else:
174 assert checkpoint_index is not None # narrowed by initial guard
175 # Negative indices count from the end, matching the legacy loader.
176 if not -len(labels) <= checkpoint_index < len(labels):
177 raise ValueError(
178 f"checkpoint_index={checkpoint_index} out of range "
179 f"[-{len(labels)}, {len(labels)}) for {model_name!r}."
180 )
181 checkpoint_value = labels[checkpoint_index]
182 return format_str.format(value=checkpoint_value)
185def list_supported_models(
186 architecture: str | None = None,
187 verified_only: bool = False,
188) -> list[str]:
189 """List all models supported by TransformerLens.
191 Args:
192 architecture: Filter by architecture ID (e.g., "GPT2LMHeadModel").
193 verified_only: If True, only return verified-to-work models.
195 Returns:
196 List of model IDs.
197 """
198 try:
199 from transformer_lens.tools.model_registry import api
201 models = api.get_supported_models(architecture=architecture, verified_only=verified_only)
202 return [m.model_id for m in models]
203 except ImportError:
204 return []
205 except Exception:
206 return []
209def check_model_support(model_id: str) -> dict:
210 """Detailed support info for a model: ``is_supported``, ``architecture_id``, ``verified``, ``suggestion``."""
211 try:
212 from transformer_lens.tools.model_registry import api
214 is_supported = api.is_model_supported(model_id)
216 if is_supported:
217 model_info = api.get_model_info(model_id)
218 return {
219 "is_supported": True,
220 "architecture_id": model_info.architecture_id,
221 "status": model_info.status,
222 "verified_date": (
223 model_info.verified_date.isoformat() if model_info.verified_date else None
224 ),
225 "suggestion": None,
226 }
227 else:
228 suggestion = api.suggest_similar_model(model_id)
229 return {
230 "is_supported": False,
231 "architecture_id": None,
232 "verified": False,
233 "verified_date": None,
234 "suggestion": suggestion,
235 }
236 except ImportError:
237 return {
238 "is_supported": None,
239 "architecture_id": None,
240 "verified": False,
241 "verified_date": None,
242 "suggestion": None,
243 "error": "Model registry not available",
244 }
245 except Exception as e:
246 return {
247 "is_supported": None,
248 "architecture_id": None,
249 "verified": False,
250 "verified_date": None,
251 "suggestion": None,
252 "error": str(e),
253 }