Coverage for transformer_lens/model_protocol.py: 100%

1"""Structural model types for the model-agnostic interpretability utilities. 

2 

3The interpretability utilities (``patching``, ``head_detector``, ``ActivationCache``) 

4only need ``cfg`` plus the ``run_with_*`` / tokenization surface and a few 

5weight-processing helpers. Typing their model parameter as this Protocol accepts 

6``TransformerBridge`` and any other structural match (e.g. ``RemoteBridge``). 

7 

8Members are typed loosely on purpose: implementations differ in signature detail 

9but agree in use. 

10""" 

11from __future__ import annotations 

12 

13from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable 

14 

15if TYPE_CHECKING: 

16 from transformer_lens.config import TransformerLensConfig 

17 

18 

19# runtime_checkable so the interp utilities' beartype-decorated signatures can 

20# isinstance-check the parameter at runtime (presence-only). 

21@runtime_checkable 

22class TransformerLensModel(Protocol): 

23 """Minimal structural interface the interpretability utilities rely on.""" 

24 

25 # Read-only property (not a bare attribute) so it is covariant: a concrete model 

26 # whose cfg is a TransformerLensConfig *subclass* (e.g. TransformerBridgeConfig) 

27 # still conforms. A mutable attribute would be invariant. 

28 @property 

29 def cfg(self) -> "TransformerLensConfig": 

30 ... 

31 

32 def run_with_cache(self, *args: Any, **kwargs: Any) -> Any: 

33 ... 

34 

35 def run_with_hooks(self, *args: Any, **kwargs: Any) -> Any: 

36 ... 

37 

38 def to_tokens(self, *args: Any, **kwargs: Any) -> Any: 

39 ... 

40 

41 

42@runtime_checkable 

43class TrainableTransformerLensModel(Protocol): 

44 """Exactly the surface the ``tools.training`` loop touches: callable with 

45 ``return_type="loss"`` plus the standard torch parameter/mode/device 

46 methods. Deliberately standalone (not extending TransformerLensModel): 

47 beartype validates protocols via ``getattr_static``, and demanding the 

48 full TL surface would spuriously reject plain ``nn.Module`` models that 

49 train fine through this loop.""" 

50 

51 def parameters(self, *args: Any, **kwargs: Any) -> Any: 

52 ... 

53 

54 def train(self, *args: Any, **kwargs: Any) -> Any: 

55 ... 

56 

57 def to(self, *args: Any, **kwargs: Any) -> Any: 

58 ... 

59 

60 def state_dict(self, *args: Any, **kwargs: Any) -> Any: 

61 ... 

62 

63 def __call__(self, *args: Any, **kwargs: Any) -> Any: 

64 ... 

65 

66 

67@runtime_checkable 

68class TransformerLensModelWithWeights(TransformerLensModel, Protocol): 

69 """Adds the weight-processing surface that ``ActivationCache``'s advanced helpers 

70 (LayerNorm folding, residual-direction projection) reach for. The bridge builds 

71 them from its adapter.""" 

72 

73 @property 

74 def blocks(self) -> Any: 

75 ... 

76 

77 @property 

78 def ln_final(self) -> Any: 

79 ... 

80 

81 def accumulated_bias(self, *args: Any, **kwargs: Any) -> Any: 

82 ... 

83 

84 def to_single_token(self, *args: Any, **kwargs: Any) -> Any: 

85 ... 

86 

87 def tokens_to_residual_directions(self, *args: Any, **kwargs: Any) -> Any: 

88 ...