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

1"""Idefics3 architecture adapter. 

2 

3Supports ``Idefics3ForConditionalGeneration`` (SmolVLM lineage — e.g. 

4ibm-granite/granite-docling-258M): a SigLIP-style vision transformer at 

5``model.vision_model``, a pixel-shuffle connector at ``model.connector``, 

6and a llama-shaped text model at ``model.text_model`` with a top-level 

7``lm_head``. 

8""" 

9 

10from typing import Any 

11 

12from transformer_lens.model_bridge.generalized_components import ( 

13 SiglipVisionEncoderBridge, 

14 VisionProjectionBridge, 

15) 

16from transformer_lens.model_bridge.supported_architectures.llama import ( 

17 LlamaArchitectureAdapter, 

18) 

19 

20 

21class Idefics3ArchitectureAdapter(LlamaArchitectureAdapter): 

22 """Architecture adapter for Idefics3ForConditionalGeneration models.""" 

23 

24 _testing_lm_attr = "model.text_model" 

25 _testing_eager = "config" 

26 

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

28 """Initialize the Idefics3 architecture adapter.""" 

29 super().__init__(cfg) 

30 

31 self.cfg.is_multimodal = True 

32 # Text model is llama-shaped (SmolLM2 in public checkpoints). 

33 self.cfg.attn_implementation = "eager" 

34 self._extract_vision_dims(cfg) 

35 self._reprefix_components("model.", "model.text_model.") 

36 

37 self.components["vision_encoder"] = SiglipVisionEncoderBridge( 

38 name="model.vision_model", config=self.cfg 

39 ) 

40 self.components["vision_projector"] = VisionProjectionBridge(name="model.connector")