Coverage for transformer_lens/model_bridge/generalized_components/pooler.py: 100%

1"""Pooler bridge component. 

2 

3This module contains the bridge component for [CLS] pooling heads. 

4""" 

5 

6from __future__ import annotations 

7 

8from transformer_lens.model_bridge.generalized_components.base import ( 

9 GeneralizedComponent, 

10) 

11 

12 

13class BertPoolerBridge(GeneralizedComponent): 

14 """Bridge component for BERT's [CLS] pooler. 

15 

16 Wraps the whole pooler, so ``hook_out`` carries the post-tanh pooled 

17 ``[CLS]`` vector rather than the pre-activation projection — the tensor 

18 ``HookedEncoder``'s ``BertPooler`` exposes as ``hook_pooler_out``, which is 

19 aliased here so code migrated from the legacy stack keeps working. 

20 """ 

21 

22 hook_aliases = {"hook_pooler_out": "hook_out"}