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

1"""Shared helpers for vision encoder bridges.""" 

2 

3from types import SimpleNamespace 

4from typing import Any 

5 

6 

7def vision_attention_config(config: Any) -> Any: 

8 """Return a config view carrying the vision tower's attention dimensions. 

9 

10 ``AttentionBridge`` reshapes q/k/v/z hooks from ``n_heads`` at fire time. A 

11 multimodal model's top-level config describes the language model, so using 

12 it directly can reshape vision activations with the text tower's head count. 

13 """ 

14 n_heads = getattr(config, "vision_num_heads", None) 

15 d_model = getattr(config, "vision_hidden_size", None) 

16 if not n_heads or not d_model: 

17 return config 

18 return SimpleNamespace(n_heads=n_heads, d_model=d_model, d_head=d_model // n_heads)