Coverage for transformer_lens/model_bridge/__init__.py: 100%

1"""Model bridge module. 

2 

3This module provides functionality to bridge between different model architectures. 

4""" 

5from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

6from transformer_lens.model_bridge.bridge import BridgeCore, RemoteBridge, TransformerBridge 

7from transformer_lens.model_bridge.component_setup import ( 

8 replace_remote_component, 

9 set_original_components, 

10 setup_blocks_bridge, 

11 setup_components, 

12 setup_submodules, 

13) 

14from transformer_lens.model_bridge.generalized_components import ( 

15 AttentionBridge, 

16 BlockBridge, 

17 EmbeddingBridge, 

18 NormalizationBridge, 

19 JointQKVAttentionBridge, 

20 JointGateUpMLPBridge, 

21 LinearBridge, 

22 MLPBridge, 

23 MoEBridge, 

24 UnembeddingBridge, 

25) 

26from transformer_lens.model_bridge.types import ( 

27 ComponentMapping, 

28 HookFunction, 

29 HookRegistry, 

30 RemoteComponent, 

31 RemoteModel, 

32 RemotePath, 

33 TransformerLensPath, 

34) 

35import transformer_lens.model_bridge.sources.transformers 

36 

37__all__ = [ 

38 "ArchitectureAdapter", 

39 "BridgeCore", 

40 "RemoteBridge", 

41 "TransformerBridge", 

42 "AttentionBridge", 

43 "BlockBridge", 

44 "EmbeddingBridge", 

45 "NormalizationBridge", 

46 "JointQKVAttentionBridge", 

47 "JointGateUpMLPBridge", 

48 "LinearBridge", 

49 "MLPBridge", 

50 "MoEBridge", 

51 "UnembeddingBridge", 

52 "replace_remote_component", 

53 "set_original_components", 

54 "setup_blocks_bridge", 

55 "setup_components", 

56 "setup_submodules", 

57 "ComponentMapping", 

58 "HookFunction", 

59 "HookRegistry", 

60 "RemoteComponent", 

61 "RemoteModel", 

62 "RemotePath", 

63 "TransformerLensPath", 

64]