transformer_lens.tools.model_registry.api module¶

Public API for the TransformerLens model registry.

This module provides a clean, programmatic interface for accessing model registry data. It supports lazy loading with in-memory caching to avoid repeated file reads.

Example usage:

>>> from transformer_lens.tools.model_registry import api  
>>> api.is_model_supported("openai-community/gpt2")  
True
>>> models = api.get_supported_models()  
>>> gpt2_models = api.get_architecture_models("GPT2LMHeadModel")  
>>> gaps = api.get_unsupported_architectures(min_models=100, top_n=10)  

transformer_lens.tools.model_registry.api.clear_cache() → None¶

Clear all cached data.

This forces data to be reloaded from disk on the next access. Useful after updating data files or for testing.

transformer_lens.tools.model_registry.api.get_all_architectures_with_stats() → list[ArchitectureStats]¶

Get statistics for all architectures (both supported and unsupported).

Returns:: List of ArchitectureStats objects for all known architectures, sorted by model count (descending)
Raises:: DataNotLoadedError – If the registry data is not available

Example

>>> stats = get_all_architectures_with_stats()  
>>> for s in stats[:5]:  
...     status = "supported" if s.is_supported else "unsupported"
...     print(f"{s.architecture_id}: {s.model_count} models ({status})")

transformer_lens.tools.model_registry.api.get_architecture_models(architecture_id: str) → list[str]¶

Get all model IDs for a given architecture.

Parameters:: architecture_id – The architecture to get models for (e.g., “GPT2LMHeadModel”)
Returns:: List of model IDs that use this architecture
Raises:: DataNotLoadedError – If the supported models data is not available

Example

>>> models = get_architecture_models("GPT2LMHeadModel")  
>>> "openai-community/gpt2" in models  
True

transformer_lens.tools.model_registry.api.get_model_architecture(model_id: str) → str | None¶

Get the architecture ID for a given model.

Parameters:: model_id – The HuggingFace model ID to look up
Returns:: The architecture ID (e.g., “GPT2LMHeadModel”), or None if not found
Raises:: DataNotLoadedError – If the supported models data is not available

Example

>>> get_model_architecture("openai-community/gpt2")  
'GPT2LMHeadModel'
>>> get_model_architecture("unknown-model")  

transformer_lens.tools.model_registry.api.get_model_info(model_id: str) → ModelEntry¶

Get full information about a specific model.

Parameters:

model_id – The HuggingFace model ID to look up

Returns:

The ModelEntry for this model

Raises:

ModelNotFoundError – If the model is not in the registry
DataNotLoadedError – If the supported models data is not available

Example

>>> info = get_model_info("openai-community/gpt2")  
>>> info.architecture_id  
'GPT2LMHeadModel'

transformer_lens.tools.model_registry.api.get_registry_stats() → dict¶

Get summary statistics about the model registry.

Returns:

total_supported_models: Number of supported models
total_supported_architectures: Number of supported architectures
total_verified: Number of verified models
total_unsupported_architectures: Number of unsupported architectures
generated_at: When the data was generated

Return type:

Dictionary with registry statistics including

Raises:

DataNotLoadedError – If the registry data is not available

Example

>>> stats = get_registry_stats()  
>>> print(f"Supported: {stats['total_supported_models']} models")  

transformer_lens.tools.model_registry.api.get_supported_architectures() → list[str]¶

Get a list of all supported architecture IDs.

Returns:: List of unique architecture IDs that TransformerLens supports
Raises:: DataNotLoadedError – If the supported models data is not available

Example

>>> archs = get_supported_architectures()  
>>> "GPT2LMHeadModel" in archs  
True

transformer_lens.tools.model_registry.api.get_supported_models(architecture: str | None = None, verified_only: bool = False) → list[ModelEntry]¶

Get a list of supported models.

Parameters:

architecture – Filter by architecture ID (e.g., “GPT2LMHeadModel”). If None, returns all supported models.
verified_only – If True, only return models that have been verified to work with TransformerLens.

Returns:

List of ModelEntry objects matching the filters

Raises:

DataNotLoadedError – If the supported models data is not available

Example

>>> models = get_supported_models(architecture="GPT2LMHeadModel")  
>>> verified = get_supported_models(verified_only=True)  

transformer_lens.tools.model_registry.api.get_unsupported_architectures(min_models: int = 0, top_n: int | None = None) → list[ArchitectureGap]¶

Get a list of unsupported architectures sorted by model count.

Parameters:

min_models – Minimum number of models for an architecture to be included. Useful for filtering out rare architectures.
top_n – Return only the top N architectures by model count. If None, returns all matching architectures.

Returns:

List of ArchitectureGap objects sorted by total_models (descending)

Raises:

DataNotLoadedError – If the architecture gaps data is not available

Example

>>> gaps = get_unsupported_architectures(min_models=100, top_n=10)  
>>> for gap in gaps:  
...     print(f"{gap.architecture_id}: {gap.total_models} models")

transformer_lens.tools.model_registry.api.is_architecture_supported(architecture_id: str) → bool¶

Check if an architecture is supported by TransformerLens.

Parameters:: architecture_id – The architecture ID to check
Returns:: True if the architecture is supported, False otherwise
Raises:: DataNotLoadedError – If the supported models data is not available

Example

>>> is_architecture_supported("GPT2LMHeadModel")  
True
>>> is_architecture_supported("SomeUnknownModel")  
False

transformer_lens.tools.model_registry.api.is_model_supported(model_id: str) → bool¶

Check if a model is supported by TransformerLens.

Parameters:: model_id – The HuggingFace model ID to check (e.g., “gpt2”, “meta-llama/Llama-2-7b-hf”)
Returns:: True if the model is in the supported models list, False otherwise
Raises:: DataNotLoadedError – If the supported models data is not available

Example

>>> is_model_supported("openai-community/gpt2")  
True
>>> is_model_supported("some-unsupported-model")  
False

transformer_lens.tools.model_registry.api.suggest_similar_model(model_id: str) → str | None¶

Suggest a similar supported model for an unsupported model ID.

This function attempts to find a supported model that is similar to the requested model based on naming patterns. Useful for providing helpful suggestions when a user tries to use an unsupported model.

Parameters:: model_id – The model ID that is not supported
Returns:: A suggested model ID, or None if no similar model is found
Raises:: DataNotLoadedError – If the supported models data is not available

Example

>>> suggest_similar_model("bigscience/bloom-560m")  
'bigscience/bloom-1b1'