Coverage for transformer_lens/model_bridge/supported_architectures/wav2vec2.py: 94%

1"""Wav2Vec2 architecture adapter. 

2 

3Wav2Vec2Model's module tree is identical to HubertModel's (zero symmetric 

4difference over named_modules), so the HuBERT adapter applies wholesale; only 

5the ForCTC nesting attribute differs (``wav2vec2.`` vs ``hubert.``). 

6""" 

7 

8from typing import Any 

9 

10from transformer_lens.model_bridge.generalized_components import UnembeddingBridge 

11from transformer_lens.model_bridge.supported_architectures.hubert import ( 

12 HubertArchitectureAdapter, 

13) 

14 

15 

16class Wav2Vec2ArchitectureAdapter(HubertArchitectureAdapter): 

17 """Adapter for Wav2Vec2Model (bare encoder) and Wav2Vec2ForCTC.""" 

18 

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

20 """Detect nesting under 'wav2vec2.' and add the CTC head when present. 

21 

22 The registered "Wav2Vec2ForPreTraining" architecture string exists so 

23 checkpoints that DECLARE that class (facebook/wav2vec2-base/-large) 

24 boot their encoder via AutoModel -> Wav2Vec2Model. Wrapping the 

25 pretraining class itself is refused: its forward returns a 

26 Wav2Vec2ForPreTrainingOutput (projected quantizer states, no 

27 last_hidden_state), which no bridge output contract fits, and the 

28 quantizer head has no interpretability surface. 

29 """ 

30 if type(hf_model).__name__ == "Wav2Vec2ForPreTraining": 

31 raise NotImplementedError( 

32 "Wav2Vec2ForPreTraining cannot be wrapped directly — boot the " 

33 "checkpoint without model_class to get its encoder " 

34 "(Wav2Vec2Model), or use Wav2Vec2ForCTC for the CTC head." 

35 ) 

36 if hasattr(hf_model, "wav2vec2"): 

37 self.component_mapping = self._build_component_mapping(prefix="wav2vec2.") 

38 if hasattr(hf_model, "lm_head"): 38 ↛ exitline 38 didn't return from function 'prepare_model' because the condition on line 38 was always true

39 self.component_mapping["unembed"] = UnembeddingBridge(name="lm_head")