Coverage for transformer_lens/__init__.py: 74%

1from . import ( 

2 conversion_utils, 

3 evals, 

4 factories, 

5 head_detector, 

6 hook_points, 

7 patching, 

8 supported_models, 

9 tools, 

10 utilities, 

11) 

12from .ActivationCache import ActivationCache 

13from .cache.key_value_cache import TransformerLensKeyValueCache 

14from .cache.key_value_cache_entry import TransformerLensKeyValueCacheEntry 

15from .config import TransformerBridgeConfig 

16from .FactoredMatrix import FactoredMatrix 

17 

18# KEPT infrastructure: HookedRootModule (with HookPoint) is the supported way 

19# to hook arbitrary nn.Modules; it was never part of the legacy model-class 

20# removal. 

21from .HookedRootModule import HookedRootModule 

22 

23# LIT integration (optional, requires lit-nlp package) 

24try: 

25 from . import lit 

26except ImportError: 

27 # LIT is an optional dependency 

28 lit = None # type: ignore 

29 

30from .SVDInterpreter import SVDInterpreter 

31 

32# Removed in 4.0: directed messages so `from transformer_lens import HookedTransformer` 

33# (and the other deleted top-level names) fail with a migration pointer instead of a 

34# bare AttributeError. Submodule-path imports (`from transformer_lens.HookedTransformer 

35# import ...`) raise ModuleNotFoundError before this hook runs and can't be intercepted here. 

36_REMOVED_IN_4_0 = { 

37 "HookedTransformer": "Use TransformerBridge.boot_transformers(name), then " 

38 "enable_compatibility_mode() for HookedTransformer-equivalent numerics.", 

39 "HookedEncoder": "Use TransformerBridge.boot_transformers(name) on a BERT model.", 

40 "HookedEncoderDecoder": "Use TransformerBridge.boot_transformers(name) on a T5 model.", 

41 "HookedAudioEncoder": "Use TransformerBridge.boot_transformers(name) on a HuBERT/Wav2Vec2 model.", 

42 "BertNextSentencePrediction": "Use TransformerBridge.boot_transformers(name, " 

43 "model_class=BertForNextSentencePrediction).predict_next_sentence(a, b).", 

44 "HookedTransformerConfig": "Use TransformerBridgeConfig.", 

45 "train": "Use transformer_lens.tools.training (train / TrainConfig).", 

46 "loading": "Model names/aliases moved to transformer_lens.supported_models; " 

47 "config derivation is now internal to TransformerBridge's adapters.", 

48 "loading_from_pretrained": "Config derivation is now internal to TransformerBridge; " 

49 "checkpoint labels live in transformer_lens.tools.model_registry.checkpoints.", 

50 "utils": "Use transformer_lens.utilities (same names).", 

51 "components": "The HookedTransformer component tree was removed; TransformerBridge " 

52 "uses transformer_lens.model_bridge.generalized_components.", 

53} 

54 

55 

56def __getattr__(name: str): 

57 # Lazy: model_bridge is import-heavy and importing it eagerly here would 

58 # risk cycles with modules the bridge itself imports. 

59 if name == "TransformerBridge": 59 ↛ 60line 59 didn't jump to line 60 because the condition on line 59 was never true

60 from .model_bridge import TransformerBridge 

61 

62 return TransformerBridge 

63 if name in _REMOVED_IN_4_0: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true

64 raise AttributeError( 

65 f"{name!r} was removed in TransformerLens 4.0. {_REMOVED_IN_4_0[name]} " 

66 "See docs/source/content/migrating_to_v4.md." 

67 ) 

68 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

69 

70 

71def __dir__(): 

72 return sorted(set(globals()) | {"TransformerBridge"}) 

73 

74 

75import os as _os # noqa: E402 

76 

77# Unconditional: without it, any model whose config writes an integral value for 

78# a float field cannot be loaded at all. See enable_hf_numeric_tower. 

79from .utilities.hf_utils import ( # noqa: E402 

80 enable_hf_numeric_tower as _enable_hf_numeric_tower, 

81) 

82 

83_enable_hf_numeric_tower() 

84 

85if _os.environ.get("TRANSFORMERLENS_HF_RETRY") == "1": 85 ↛ 90line 85 didn't jump to line 90 because the condition on line 85 was always true

86 from .utilities.hf_utils import enable_hf_retry as _enable_hf_retry # noqa: E402 

87 

88 _enable_hf_retry() 

89 

90__all__ = [ 

91 "TransformerBridge", 

92 "TransformerBridgeConfig", 

93 "FactoredMatrix", 

94 "ActivationCache", 

95 "SVDInterpreter", 

96 "HookedRootModule", 

97 "TransformerLensKeyValueCache", 

98 "TransformerLensKeyValueCacheEntry", 

99 "conversion_utils", 

100 "factories", 

101 "utilities", 

102 "tools", 

103]