Coverage for transformer_lens/__init__.py: 80%

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 error. We raise ImportError, not AttributeError: the from-import machinery calls 

35# getattr and, on AttributeError, discards it and raises its own unchained "cannot import 

36# name" ImportError, so our pointer would be lost for exactly the spelling downstream 

37# packages use. ImportError reaches both `import transformer_lens; transformer_lens.X` and 

38# `from transformer_lens import X`. Trade-off: `hasattr(transformer_lens, "X")` now raises 

39# instead of returning False, and `getattr(transformer_lens, "X", default)` raises instead 

40# of returning the default (both only swallow AttributeError) — intended, so feature- 

41# detection code gets the migration pointer too. A subclass of both is impossible: 

42# ImportError's C layout conflicts with AttributeError's. Submodule-path imports 

43# (`from transformer_lens.HookedTransformer import ...`) raise ModuleNotFoundError before 

44# this hook runs and can't be intercepted here. 

45_REMOVED_IN_4_0 = { 

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

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

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

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

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

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

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

53 "HookedTransformerConfig": "Use TransformerBridgeConfig.", 

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

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

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

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

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

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

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

61 "uses transformer_lens.model_bridge.generalized_components.", 

62} 

63 

64 

65def __getattr__(name: str): 

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

67 # risk cycles with modules the bridge itself imports. 

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

69 from .model_bridge import TransformerBridge 

70 

71 return TransformerBridge 

72 if name in _REMOVED_IN_4_0: 

73 raise ImportError( 

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

75 "See docs/source/content/migrating_to_v4.md.", 

76 name=name, 

77 ) 

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

79 

80 

81def __dir__(): 

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

83 

84 

85import os as _os # noqa: E402 

86 

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

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

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

90 enable_hf_numeric_tower as _enable_hf_numeric_tower, 

91) 

92 

93_enable_hf_numeric_tower() 

94 

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

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

97 

98 _enable_hf_retry() 

99 

100__all__ = [ 

101 "TransformerBridge", 

102 "TransformerBridgeConfig", 

103 "FactoredMatrix", 

104 "ActivationCache", 

105 "SVDInterpreter", 

106 "HookedRootModule", 

107 "TransformerLensKeyValueCache", 

108 "TransformerLensKeyValueCacheEntry", 

109 "conversion_utils", 

110 "factories", 

111 "utilities", 

112 "tools", 

113]