Coverage for transformer_lens/utilities/position_ids.py: 100%

1"""Shared eligibility gate for mask-derived position IDs.""" 

2 

3import inspect 

4 

5from torch import nn 

6 

7 

8def accepts_mask_derived_position_ids(model: nn.Module) -> bool: 

9 """Accept explicit/kwargs positions unless the model owns their derivation.""" 

10 parameters = inspect.signature(model.forward).parameters 

11 if "position_ids" not in parameters and not any( 

12 p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values() 

13 ): 

14 return False 

15 # ``get_rope_index`` lives on the inner text model, not the ForConditionalGeneration wrapper. 

16 for module in (model, getattr(model, "model", None), getattr(model, "language_model", None)): 

17 if module is not None and hasattr(module, "get_rope_index"): 

18 return False 

19 # Config-level backstop for mRoPE models that spell the derivation differently; the 

20 # section list is what makes positions 3-D. 

21 config = getattr(model, "config", None) 

22 for candidate in (config, getattr(config, "text_config", None)): 

23 scaling = getattr(candidate, "rope_scaling", None) 

24 if isinstance(scaling, dict) and "mrope_section" in scaling: 

25 return False 

26 # A positional embedding that takes the mask derives positions for itself; only 

27 # embeddings that override nn.Embedding.forward are worth inspecting. 

28 for module in model.modules(): 

29 if isinstance(module, nn.Embedding) and type(module).forward is not nn.Embedding.forward: 

30 if "attention_mask" in inspect.signature(module.forward).parameters: 

31 return False 

32 return True