Coverage for transformer_lens/tools/model_registry/__init__.py: 100%

1"""Model Registry tools for TransformerLens. 

2 

3This package provides tools for discovering and documenting HuggingFace models 

4that are compatible with TransformerLens. 

5 

6Main modules: 

7 - api: Public API for programmatic access to model registry data 

8 - schemas: Data classes for model entries, architecture gaps, etc. 

9 - verification: Verification tracking for model compatibility 

10 - exceptions: Custom exceptions for the model registry 

11 

12Example usage: 

13 >>> from transformer_lens.tools.model_registry import api # doctest: +SKIP 

14 >>> api.is_model_supported("openai-community/gpt2") # doctest: +SKIP 

15 True 

16 >>> models = api.get_architecture_models("GPT2LMHeadModel") # doctest: +SKIP 

17""" 

18 

19from .exceptions import ( 

20 ArchitectureNotSupportedError, 

21 DataNotLoadedError, 

22 DataValidationError, 

23 ModelNotFoundError, 

24 ModelRegistryError, 

25) 

26from .schemas import ( 

27 ArchitectureAnalysis, 

28 ArchitectureGap, 

29 ArchitectureGapsReport, 

30 ArchitectureStats, 

31 ModelEntry, 

32 ModelMetadata, 

33 ScanInfo, 

34 SupportedModelsReport, 

35) 

36from .verification import VerificationHistory, VerificationRecord 

37 

38# Canonical set of HuggingFace architecture class names supported by TransformerBridge. 

39# These must match the exact strings found in HF model config.architectures[] 

40# and correspond to adapters registered in architecture_adapter_factory.py. 

41# 

42# Internal-only architectures (NanoGPT, MinGPT, NeelSoluOld, GPT2LMHeadCustomModel) 

43# are excluded since they never appear on HuggingFace Hub. 

44HF_SUPPORTED_ARCHITECTURES: set[str] = { 

45 "ApertusForCausalLM", 

46 "BertForMaskedLM", 

47 "BloomForCausalLM", 

48 "CodeGenForCausalLM", 

49 "CohereForCausalLM", 

50 "DeepseekV3ForCausalLM", 

51 "FalconForCausalLM", 

52 "GemmaForCausalLM", 

53 "Gemma2ForCausalLM", 

54 "Gemma3ForCausalLM", 

55 "Gemma3ForConditionalGeneration", 

56 "GraniteForCausalLM", 

57 "GraniteMoeForCausalLM", 

58 "GraniteMoeHybridForCausalLM", 

59 "GPT2LMHeadModel", 

60 "GPTBigCodeForCausalLM", 

61 "GptOssForCausalLM", 

62 "GPTJForCausalLM", 

63 "GPTNeoForCausalLM", 

64 "OpenELMForCausalLM", 

65 "GPTNeoXForCausalLM", 

66 "HubertForCTC", 

67 "HubertModel", 

68 "InternLM2ForCausalLM", 

69 "LlamaForCausalLM", 

70 "LlavaForConditionalGeneration", 

71 "LlavaNextForConditionalGeneration", 

72 "LlavaOnevisionForConditionalGeneration", 

73 "MambaForCausalLM", 

74 "Mamba2ForCausalLM", 

75 "MPTForCausalLM", 

76 "MistralForCausalLM", 

77 "MixtralForCausalLM", 

78 "Olmo2ForCausalLM", 

79 "Olmo3ForCausalLM", 

80 "OlmoForCausalLM", 

81 "OlmoeForCausalLM", 

82 "OPTForCausalLM", 

83 "PhiForCausalLM", 

84 "Phi3ForCausalLM", 

85 "QwenForCausalLM", 

86 "Qwen2ForCausalLM", 

87 "Qwen3ForCausalLM", 

88 "Qwen3NextForCausalLM", 

89 "Qwen3_5ForCausalLM", 

90 "StableLmForCausalLM", 

91 "T5ForConditionalGeneration", 

92 "XGLMForCausalLM", 

93} 

94 

95__all__ = [ 

96 # Constants 

97 "HF_SUPPORTED_ARCHITECTURES", 

98 # Exceptions 

99 "ModelRegistryError", 

100 "ModelNotFoundError", 

101 "ArchitectureNotSupportedError", 

102 "DataNotLoadedError", 

103 "DataValidationError", 

104 # Schemas 

105 "ModelEntry", 

106 "ModelMetadata", 

107 "ScanInfo", 

108 "ArchitectureGap", 

109 "ArchitectureStats", 

110 "ArchitectureAnalysis", 

111 "SupportedModelsReport", 

112 "ArchitectureGapsReport", 

113 # Verification 

114 "VerificationRecord", 

115 "VerificationHistory", 

116]