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 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 "TransformerBridge", 

40 "AttentionBridge", 

41 "BlockBridge", 

42 "EmbeddingBridge", 

43 "NormalizationBridge", 

44 "JointQKVAttentionBridge", 

45 "JointGateUpMLPBridge", 

46 "LinearBridge", 

47 "MLPBridge", 

48 "MoEBridge", 

49 "UnembeddingBridge", 

50 "replace_remote_component", 

51 "set_original_components", 

52 "setup_blocks_bridge", 

53 "setup_components", 

54 "setup_submodules", 

55 "ComponentMapping", 

56 "HookFunction", 

57 "HookRegistry", 

58 "RemoteComponent", 

59 "RemoteModel", 

60 "RemotePath", 

61 "TransformerLensPath", 

62]