Coverage for transformer_lens/tools/model_registry/api.py: 60%
144 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Public API for the TransformerLens model registry.
3This module provides a clean, programmatic interface for accessing model registry
4data. It supports lazy loading with in-memory caching to avoid repeated file reads.
6Example usage:
7 >>> from transformer_lens.tools.model_registry import api # doctest: +SKIP
8 >>> api.is_model_supported("openai-community/gpt2") # doctest: +SKIP
9 True
10 >>> models = api.get_supported_models() # doctest: +SKIP
11 >>> gpt2_models = api.get_architecture_models("GPT2LMHeadModel") # doctest: +SKIP
12 >>> gaps = api.get_unsupported_architectures(min_models=100, top_n=10) # doctest: +SKIP
13"""
15import json
16import logging
17from pathlib import Path
18from threading import Lock
19from typing import Optional
21from . import ARCHITECTURE_ALIASES
22from .exceptions import DataNotLoadedError, ModelNotFoundError
23from .schemas import (
24 ArchitectureGap,
25 ArchitectureGapsReport,
26 ArchitectureStats,
27 ModelEntry,
28 SupportedModelsReport,
29)
30from .verification import VerificationHistory
32logger = logging.getLogger(__name__)
34# Module-level cache for lazy loading
35_cache: dict[str, object] = {}
36_cache_lock = Lock()
38# Default data directory (relative to this module)
39_DATA_DIR = Path(__file__).parent / "data"
42def _load_json(filename: str) -> dict:
43 """Load a JSON file from the data directory.
45 Args:
46 filename: Name of the JSON file
48 Returns:
49 Parsed JSON data as a dictionary
51 Raises:
52 DataNotLoadedError: If the file doesn't exist or can't be read
53 """
54 path = _DATA_DIR / filename
55 if not path.exists(): 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true
56 raise DataNotLoadedError(filename, str(path))
57 try:
58 with open(path) as f:
59 return json.load(f)
60 except json.JSONDecodeError as e:
61 raise DataNotLoadedError(filename, str(path)) from e
64def _get_supported_models_report() -> SupportedModelsReport:
65 """Get the cached supported models report, loading if necessary.
67 Returns:
68 The SupportedModelsReport instance
70 Raises:
71 DataNotLoadedError: If the data files are not available
72 """
73 cache_key = "supported_models"
74 with _cache_lock:
75 if cache_key not in _cache:
76 data = _load_json("supported_models.json")
77 _cache[cache_key] = SupportedModelsReport.from_dict(data)
78 result = _cache[cache_key]
79 assert isinstance(result, SupportedModelsReport)
80 return result
83def _get_architecture_gaps_report() -> ArchitectureGapsReport:
84 """Get the cached architecture gaps report, loading if necessary.
86 Returns:
87 The ArchitectureGapsReport instance
89 Raises:
90 DataNotLoadedError: If the data file is not available
91 """
92 cache_key = "architecture_gaps"
93 with _cache_lock:
94 if cache_key not in _cache: 94 ↛ 97line 94 didn't jump to line 97 because the condition on line 94 was always true
95 data = _load_json("architecture_gaps.json")
96 _cache[cache_key] = ArchitectureGapsReport.from_dict(data)
97 result = _cache[cache_key]
98 assert isinstance(result, ArchitectureGapsReport)
99 return result
102def _get_verification_history() -> VerificationHistory:
103 """Get the cached verification history, loading if necessary.
105 Returns:
106 The VerificationHistory instance
108 Raises:
109 DataNotLoadedError: If the data files are not available
110 """
111 cache_key = "verification_history"
112 with _cache_lock:
113 if cache_key not in _cache:
114 data = _load_json("verification_history.json")
115 _cache[cache_key] = VerificationHistory.from_dict(data)
116 result = _cache[cache_key]
117 assert isinstance(result, VerificationHistory)
118 return result
121def clear_cache() -> None:
122 """Clear all cached data.
124 This forces data to be reloaded from disk on the next access.
125 Useful after updating data files or for testing.
126 """
127 with _cache_lock:
128 _cache.clear()
129 logger.debug("Model registry cache cleared")
132def get_supported_models(
133 architecture: Optional[str] = None,
134 verified_only: bool = False,
135) -> list[ModelEntry]:
136 """Get a list of supported models.
138 Args:
139 architecture: Filter by architecture ID (e.g., "GPT2LMHeadModel").
140 If None, returns all supported models.
141 verified_only: If True, only return models that have been verified
142 to work with TransformerLens.
144 Returns:
145 List of ModelEntry objects matching the filters
147 Raises:
148 DataNotLoadedError: If the supported models data is not available
150 Example:
151 >>> models = get_supported_models(architecture="GPT2LMHeadModel") # doctest: +SKIP
152 >>> verified = get_supported_models(verified_only=True) # doctest: +SKIP
153 """
154 report = _get_supported_models_report()
155 models = report.models
157 if architecture:
158 models = [m for m in models if m.architecture_id == architecture]
160 if verified_only:
161 models = [m for m in models if m.status == 1]
163 return models
166def get_unsupported_architectures(
167 min_models: int = 0,
168 top_n: Optional[int] = None,
169) -> list[ArchitectureGap]:
170 """Get a list of unsupported architectures sorted by model count.
172 Args:
173 min_models: Minimum number of models for an architecture to be included.
174 Useful for filtering out rare architectures.
175 top_n: Return only the top N architectures by model count.
176 If None, returns all matching architectures.
178 Returns:
179 List of ArchitectureGap objects sorted by total_models (descending)
181 Raises:
182 DataNotLoadedError: If the architecture gaps data is not available
184 Example:
185 >>> gaps = get_unsupported_architectures(min_models=100, top_n=10) # doctest: +SKIP
186 >>> for gap in gaps: # doctest: +SKIP
187 ... print(f"{gap.architecture_id}: {gap.total_models} models")
188 """
189 report = _get_architecture_gaps_report()
190 gaps = report.gaps
192 if min_models > 0:
193 gaps = [g for g in gaps if g.total_models >= min_models]
195 # Already sorted by total_models descending in the report
196 if top_n is not None:
197 gaps = gaps[:top_n]
199 return gaps
202def is_model_supported(model_id: str) -> bool:
203 """Check if a model is supported by TransformerLens.
205 Args:
206 model_id: The HuggingFace model ID to check (e.g., "gpt2", "meta-llama/Llama-2-7b-hf")
208 Returns:
209 True if the model is in the supported models list, False otherwise
211 Raises:
212 DataNotLoadedError: If the supported models data is not available
214 Example:
215 >>> is_model_supported("openai-community/gpt2") # doctest: +SKIP
216 True
217 >>> is_model_supported("some-unsupported-model") # doctest: +SKIP
218 False
219 """
220 report = _get_supported_models_report()
221 return any(m.model_id == model_id for m in report.models)
224def get_model_architecture(model_id: str) -> Optional[str]:
225 """Get the architecture ID for a given model.
227 Args:
228 model_id: The HuggingFace model ID to look up
230 Returns:
231 The architecture ID (e.g., "GPT2LMHeadModel"), or None if not found
233 Raises:
234 DataNotLoadedError: If the supported models data is not available
236 Example:
237 >>> get_model_architecture("openai-community/gpt2") # doctest: +SKIP
238 'GPT2LMHeadModel'
239 >>> get_model_architecture("unknown-model") # doctest: +SKIP
240 """
241 report = _get_supported_models_report()
242 for model in report.models:
243 if model.model_id == model_id:
244 return model.architecture_id
245 return None
248def get_architecture_models(architecture_id: str) -> list[str]:
249 """Get all model IDs for a given architecture.
251 Args:
252 architecture_id: The architecture to get models for (e.g., "GPT2LMHeadModel")
254 Returns:
255 List of model IDs that use this architecture
257 Raises:
258 DataNotLoadedError: If the supported models data is not available
260 Example:
261 >>> models = get_architecture_models("GPT2LMHeadModel") # doctest: +SKIP
262 >>> "openai-community/gpt2" in models # doctest: +SKIP
263 True
264 """
265 report = _get_supported_models_report()
266 return [m.model_id for m in report.models if m.architecture_id == architecture_id]
269def suggest_similar_model(model_id: str) -> Optional[str]:
270 """Suggest a similar supported model for an unsupported model ID.
272 This function attempts to find a supported model that is similar to the
273 requested model based on naming patterns. Useful for providing helpful
274 suggestions when a user tries to use an unsupported model.
276 Args:
277 model_id: The model ID that is not supported
279 Returns:
280 A suggested model ID, or None if no similar model is found
282 Raises:
283 DataNotLoadedError: If the supported models data is not available
285 Example:
286 >>> suggest_similar_model("bigscience/bloom-560m") # doctest: +SKIP
287 'bigscience/bloom-1b1'
288 """
289 report = _get_supported_models_report()
291 # If the model is already supported, return None (no suggestion needed)
292 if any(m.model_id == model_id for m in report.models): 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true
293 return None
295 # Extract potential matching criteria from the model ID
296 model_id_lower = model_id.lower()
297 parts = model_id.replace("/", "-").replace("_", "-").lower().split("-")
299 # Build a scoring function for similarity
300 def score_model(candidate: ModelEntry) -> int:
301 candidate_lower = candidate.model_id.lower()
302 score = 0
304 # Same organization prefix
305 if "/" in model_id and "/" in candidate.model_id: 305 ↛ 310line 305 didn't jump to line 310 because the condition on line 305 was always true
306 if model_id.split("/")[0].lower() == candidate.model_id.split("/")[0].lower(): 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true
307 score += 10
309 # Matching parts
310 for part in parts:
311 if len(part) > 2 and part in candidate_lower: 311 ↛ 312line 311 didn't jump to line 312 because the condition on line 311 was never true
312 score += 5
314 # Architecture name hints
315 arch_hints = ["gpt", "llama", "bloom", "opt", "mistral", "gemma", "phi", "qwen"]
316 for hint in arch_hints:
317 if hint in model_id_lower and hint in candidate_lower: 317 ↛ 318line 317 didn't jump to line 318 because the condition on line 317 was never true
318 score += 8
320 return score
322 # Score all models and find the best match
323 scored = [(m, score_model(m)) for m in report.models]
324 scored = [(m, s) for m, s in scored if s > 0] # Only consider matches with some score
325 scored.sort(key=lambda x: x[1], reverse=True)
327 if scored: 327 ↛ 328line 327 didn't jump to line 328 because the condition on line 327 was never true
328 return scored[0][0].model_id
329 return None
332def get_model_info(model_id: str) -> ModelEntry:
333 """Get full information about a specific model.
335 Args:
336 model_id: The HuggingFace model ID to look up
338 Returns:
339 The ModelEntry for this model
341 Raises:
342 ModelNotFoundError: If the model is not in the registry
343 DataNotLoadedError: If the supported models data is not available
345 Example:
346 >>> info = get_model_info("openai-community/gpt2") # doctest: +SKIP
347 >>> info.architecture_id # doctest: +SKIP
348 'GPT2LMHeadModel'
349 """
350 report = _get_supported_models_report()
351 for model in report.models:
352 if model.model_id == model_id: 352 ↛ 353line 352 didn't jump to line 353 because the condition on line 352 was never true
353 return model
355 # Model not found - try to suggest an alternative
356 suggestion = suggest_similar_model(model_id)
357 raise ModelNotFoundError(model_id, suggestion)
360def get_supported_architectures() -> list[str]:
361 """Get a list of all supported architecture IDs.
363 Returns:
364 List of unique architecture IDs that TransformerLens supports
366 Raises:
367 DataNotLoadedError: If the supported models data is not available
369 Example:
370 >>> archs = get_supported_architectures() # doctest: +SKIP
371 >>> "GPT2LMHeadModel" in archs # doctest: +SKIP
372 True
373 """
374 report = _get_supported_models_report()
375 return list(sorted(set(m.architecture_id for m in report.models)))
378def get_all_architectures_with_stats() -> list[ArchitectureStats]:
379 """Get statistics for all architectures (both supported and unsupported).
381 Returns:
382 List of ArchitectureStats objects for all known architectures,
383 sorted by model count (descending)
385 Raises:
386 DataNotLoadedError: If the registry data is not available
388 Example:
389 >>> stats = get_all_architectures_with_stats() # doctest: +SKIP
390 >>> for s in stats[:5]: # doctest: +SKIP
391 ... status = "supported" if s.is_supported else "unsupported"
392 ... print(f"{s.architecture_id}: {s.model_count} models ({status})")
393 """
394 gaps_report = _get_architecture_gaps_report()
395 supported_report = _get_supported_models_report()
397 # Build stats for supported architectures
398 arch_stats: dict[str, ArchitectureStats] = {}
399 for model in supported_report.models:
400 arch_id = model.architecture_id
401 if arch_id not in arch_stats:
402 arch_stats[arch_id] = ArchitectureStats(
403 architecture_id=arch_id,
404 is_supported=True,
405 model_count=0,
406 verified_count=0,
407 example_models=[],
408 )
409 stats_obj = arch_stats[arch_id]
410 stats_obj.model_count += 1
411 if model.status == 1:
412 stats_obj.verified_count += 1
413 if len(stats_obj.example_models) < 5:
414 stats_obj.example_models.append(model.model_id)
416 # Add stats for unsupported architectures
417 for gap in gaps_report.gaps:
418 if gap.architecture_id not in arch_stats:
419 arch_stats[gap.architecture_id] = ArchitectureStats(
420 architecture_id=gap.architecture_id,
421 is_supported=False,
422 model_count=gap.total_models,
423 verified_count=0,
424 example_models=[],
425 )
427 result = sorted(arch_stats.values(), key=lambda x: x.model_count, reverse=True)
428 return result
431def is_architecture_supported(architecture_id: str) -> bool:
432 """Check if an architecture is supported by TransformerLens.
434 Args:
435 architecture_id: The architecture ID to check
437 Returns:
438 True if the architecture is supported, False otherwise
440 Raises:
441 DataNotLoadedError: If the supported models data is not available
443 Example:
444 >>> is_architecture_supported("GPT2LMHeadModel") # doctest: +SKIP
445 True
446 >>> is_architecture_supported("SomeUnknownModel") # doctest: +SKIP
447 False
448 """
449 # Hub configs may report an alias casing (e.g. JetMoEForCausalLM) while rows
450 # are keyed by the transformers class name; normalize before scanning.
451 architecture_id = ARCHITECTURE_ALIASES.get(architecture_id, architecture_id)
452 report = _get_supported_models_report()
453 return any(m.architecture_id == architecture_id for m in report.models)
456def get_registry_stats() -> dict:
457 """Get summary statistics about the model registry.
459 Returns:
460 Dictionary with registry statistics including:
461 - total_supported_models: Number of supported models
462 - total_supported_architectures: Number of supported architectures
463 - total_verified: Number of verified models
464 - total_provisional: Number of provisional (structural-only) models
465 - total_unsupported_architectures: Number of unsupported architectures
466 - generated_at: When the data was generated
468 Raises:
469 DataNotLoadedError: If the registry data is not available
471 Example:
472 >>> stats = get_registry_stats() # doctest: +SKIP
473 >>> print(f"Supported: {stats['total_supported_models']} models") # doctest: +SKIP
474 """
475 supported = _get_supported_models_report()
476 gaps = _get_architecture_gaps_report()
478 return {
479 "total_supported_models": supported.total_models,
480 "total_supported_architectures": supported.total_architectures,
481 "total_verified": supported.total_verified,
482 "total_provisional": supported.total_provisional,
483 "total_unsupported_architectures": gaps.total_unsupported_architectures,
484 "total_unsupported_models": gaps.total_unsupported_models,
485 "supported_generated_at": supported.generated_at.isoformat(),
486 "gaps_generated_at": gaps.generated_at.isoformat(),
487 }