Coverage for transformer_lens/model_bridge/supported_architectures/roberta.py: 100%

1"""RoBERTa architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.model_bridge.generalized_components import ( 

6 AttentionBridge, 

7 BlockBridge, 

8 EmbeddingBridge, 

9 LinearBridge, 

10 MLPBridge, 

11 NormalizationBridge, 

12 PosEmbedBridge, 

13 UnembeddingBridge, 

14) 

15from transformer_lens.model_bridge.supported_architectures.bert import ( 

16 BertArchitectureAdapter, 

17) 

18 

19 

20class RobertaArchitectureAdapter(BertArchitectureAdapter): 

21 """Architecture adapter for ``RobertaForMaskedLM``: BERT's layer, MLM head under ``lm_head.*``. 

22 

23 Positions start at ``padding_idx + 1 = 2``, so ``W_pos`` drops the table's first two rows. 

24 """ 

25 

26 def __init__(self, cfg: Any) -> None: 

27 """Initialize the RoBERTa architecture adapter.""" 

28 super().__init__(cfg) 

29 

30 self.component_mapping = { 

31 "embed": EmbeddingBridge(name="roberta.embeddings.word_embeddings"), 

32 # Real [1, d_model] table at type_vocab_size = 1; HF adds it on every forward. 

33 "token_type_embed": EmbeddingBridge(name="roberta.embeddings.token_type_embeddings"), 

34 "pos_embed": PosEmbedBridge(name="roberta.embeddings.position_embeddings"), 

35 "embed_ln": NormalizationBridge( 

36 name="roberta.embeddings.LayerNorm", 

37 config=self.cfg, 

38 use_native_layernorm_autograd=True, 

39 ), 

40 "blocks": BlockBridge( 

41 name="roberta.encoder.layer", 

42 # Same split intermediate.dense / output.dense layout as BertLayer, see bert.py. 

43 hook_alias_overrides={ 

44 "hook_mlp_out": "mlp.out.hook_out", 

45 "hook_mlp_in": "mlp.in.hook_in", 

46 }, 

47 submodules={ 

48 "ln1": NormalizationBridge( 

49 name="attention.output.LayerNorm", 

50 config=self.cfg, 

51 use_native_layernorm_autograd=True, 

52 ), 

53 "ln2": NormalizationBridge( 

54 name="output.LayerNorm", 

55 config=self.cfg, 

56 use_native_layernorm_autograd=True, 

57 ), 

58 "attn": AttentionBridge( 

59 name="attention", 

60 config=self.cfg, 

61 submodules={ 

62 "q": LinearBridge(name="self.query"), 

63 "k": LinearBridge(name="self.key"), 

64 "v": LinearBridge(name="self.value"), 

65 "o": LinearBridge(name="output.dense"), 

66 }, 

67 ), 

68 "mlp": MLPBridge( 

69 name=None, 

70 config=self.cfg, 

71 submodules={ 

72 "in": LinearBridge(name="intermediate.dense"), 

73 "out": LinearBridge(name="output.dense"), 

74 }, 

75 ), 

76 }, 

77 ), 

78 "mlm_head": LinearBridge(name="lm_head.dense"), 

79 "unembed": UnembeddingBridge(name="lm_head.decoder"), 

80 "ln_final": NormalizationBridge( 

81 name="lm_head.layer_norm", 

82 config=self.cfg, 

83 use_native_layernorm_autograd=True, 

84 ), 

85 } 

86 

87 def prepare_model(self, hf_model: Any) -> None: 

88 """Read the usable context off the loaded position table.""" 

89 super().prepare_model(hf_model) 

90 # Drop the padding_idx + 1 reserved rows; n_ctx + 2 rows is what makes W_pos slice. 

91 table = hf_model.roberta.embeddings.position_embeddings 

92 self.cfg.n_ctx = table.num_embeddings - ((table.padding_idx or 0) + 1)