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

1"""TL-native model source for TransformerBridge.""" 

2from __future__ import annotations 

3 

4import copy as _copy 

5from typing import Any, Optional, Union, overload 

6 

7import torch 

8 

9from transformer_lens.config import TransformerBridgeConfig 

10from transformer_lens.model_bridge.bridge import TransformerBridge 

11from transformer_lens.model_bridge.sources._bridge_builder import ( 

12 build_bridge_from_module, 

13) 

14from transformer_lens.model_bridge.sources.native.init import initialize_native_model 

15from transformer_lens.model_bridge.sources.native.model import ( 

16 NativeAttention, 

17 NativeBlock, 

18 NativeMLP, 

19 NativeModel, 

20) 

21 

22 

23@overload 

24def boot( 

25 config: TransformerBridgeConfig, 

26 tokenizer: Optional[Any] = None, 

27 device: Optional[Union[str, torch.device]] = None, 

28 dtype: Optional[torch.dtype] = None, 

29 model_name: str = "native", 

30) -> TransformerBridge: 

31 ... 

32 

33 

34@overload 

35def boot( 

36 config: dict[str, Any], 

37 tokenizer: Optional[Any] = None, 

38 device: Optional[Union[str, torch.device]] = None, 

39 dtype: Optional[torch.dtype] = None, 

40 model_name: str = "native", 

41) -> TransformerBridge: 

42 ... 

43 

44 

45def boot( 

46 config: Any, 

47 tokenizer: Optional[Any] = None, 

48 device: Optional[Union[str, torch.device]] = None, 

49 dtype: Optional[torch.dtype] = None, 

50 model_name: str = "native", 

51) -> TransformerBridge: 

52 """Build a bridge around a small, randomly-initialized TL-native model. 

53 

54 No HuggingFace Hub call, no ``transformers`` import. ``config.init_mode`` 

55 and ``config.seed`` control reproducibility. 

56 """ 

57 if not isinstance(config, (TransformerBridgeConfig, dict)): 

58 raise TypeError( 

59 "boot_native expected a TransformerBridgeConfig or dict, " 

60 f"got {type(config).__name__}. Construct a TransformerBridgeConfig " 

61 "with the same fields; legacy config classes are deprecated and " 

62 "are not accepted." 

63 ) 

64 

65 cfg: TransformerBridgeConfig 

66 if isinstance(config, dict): 

67 cfg = TransformerBridgeConfig.from_dict(config) 

68 else: 

69 # Deep-copy so NativeModel's default-resolution writes don't land 

70 # on the caller's config. 

71 cfg = _copy.deepcopy(config) 

72 

73 # Foreign architecture strings would dispatch to the wrong adapter and 

74 # crash deep in prepare_model. Refuse them with a pointing message. 

75 if cfg.architecture not in (None, "TransformerLensNative"): 

76 raise ValueError( 

77 f"boot_native cannot build a {cfg.architecture!r} model — " 

78 f"it only constructs the TL-native architecture. Either clear " 

79 f"config.architecture or set it to 'TransformerLensNative', " 

80 f"or use boot_transformers / build_bridge_from_module for " 

81 f"non-native architectures." 

82 ) 

83 architecture = "TransformerLensNative" 

84 

85 # Fork RNG around construction + init when seeded so neither nn.Linear's 

86 # default reset_parameters nor our scoped init perturb the caller's RNG. 

87 # When custom init is disabled, construction keeps PyTorch's normal global 

88 # RNG semantics and cfg.seed has no initialization work to control. 

89 if cfg.init_weights and cfg.seed is not None: 

90 with torch.random.fork_rng(devices=[]): 

91 model = NativeModel(cfg) 

92 initialize_native_model(model, cfg) 

93 else: 

94 model = NativeModel(cfg) 

95 if cfg.init_weights: 

96 initialize_native_model(model, cfg) 

97 

98 if device is not None: 

99 model = model.to(device) 

100 if dtype is not None: 

101 model = model.to(dtype=dtype) 

102 

103 return build_bridge_from_module( 

104 model, 

105 architecture=architecture, 

106 tl_config=cfg, 

107 tokenizer=tokenizer, 

108 dtype=dtype, 

109 device=device, 

110 model_name=model_name, 

111 ) 

112 

113 

114# Attach to TransformerBridge as a staticmethod, matching boot_transformers / boot_vllm. 

115setattr(TransformerBridge, "boot_native", staticmethod(boot)) 

116 

117 

118__all__ = [ 

119 "NativeAttention", 

120 "NativeBlock", 

121 "NativeMLP", 

122 "NativeModel", 

123 "boot", 

124 "initialize_native_model", 

125]