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 # fairseq-style (RoBERTa) embeddings start positions at padding_idx + 1; an injected 

27 # 0-based arange would put content tokens on the reserved rows. 

28 for module in model.modules(): 

29 if callable(getattr(module, "create_position_ids_from_input_ids", None)): 

30 return False 

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

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

33 for module in model.modules(): 

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

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

36 return False 

37 return True