Coverage for transformer_lens/model_bridge/sources/_driver_base.py: 100%

1"""Optional base class for :class:`Driver` implementations. 

2 

3The protocol is duck-typed; inheriting is convenient, not required. 

4""" 

5from __future__ import annotations 

6 

7from abc import ABC, abstractmethod 

8from typing import Any, Mapping 

9 

10from transformer_lens.config import TransformerBridgeConfig 

11from transformer_lens.model_bridge.driver_protocol import ( 

12 ForwardResult, 

13 Intervention, 

14 TensorLike, 

15) 

16 

17 

18class DriverBase(ABC): 

19 """Defaults for the optional Driver members. Subclasses implement ``forward`` 

20 and override the rest only when they can do better.""" 

21 

22 architecture: str = "" 

23 # The bridge overwrites this slot at construction (registry − non_fireable) 

24 # when it's empty. Whitelist-semantic drivers (e.g. Inspect) declare a 

25 # non-empty set in the subclass to keep it. 

26 supported_hook_points: frozenset[str] = frozenset() 

27 non_fireable_hook_points: frozenset[str] = frozenset() 

28 

29 # Subclasses override with the capability strings they actually serve. 

30 _supported_features: frozenset[str] = frozenset() 

31 

32 # True if forward() returns logits for every position, so loss is computable. 

33 # Drivers that synthesize only the final position set this False; the bridge 

34 # then refuses return_type=loss/both rather than return nan. 

35 provides_sequence_logits: bool = True 

36 

37 def __init__( 

38 self, 

39 bridge_config: TransformerBridgeConfig, 

40 tokenizer: Any, 

41 *, 

42 architecture: str | None = None, 

43 ) -> None: 

44 self.bridge_config = bridge_config 

45 self.tokenizer = tokenizer 

46 # Resolution order: explicit kwarg > bridge_config field > class default. 

47 self.architecture = ( 

48 architecture or getattr(bridge_config, "architecture", "") or self.architecture 

49 ) 

50 

51 @abstractmethod 

52 def forward( 

53 self, 

54 input_ids: TensorLike | None = None, 

55 *, 

56 capture: tuple[str, ...] = (), 

57 intervene: Mapping[str, Intervention] | None = None, 

58 max_new_tokens: int = 1, 

59 return_logits: bool = True, 

60 **kwargs: Any, 

61 ) -> ForwardResult: 

62 ... 

63 

64 def close(self) -> None: 

65 """No-op default. Override when the driver owns releasable resources.""" 

66 

67 def supports(self, feature: str) -> bool: 

68 return feature in self._supported_features 

69 

70 # Torch-specific surface (parameters, named_parameters, state_dict, weight 

71 # access) is NOT defined here. Drivers that serve those methods provide 

72 # them directly; callers route via supports("...") + hasattr/getattr.