Coverage for transformer_lens/tools/analysis/_model_state.py: 100%

1"""Shared model-state and argument validation for analysis tools.""" 

2 

3from numbers import Integral 

4from typing import Any, List, Optional 

5 

6import torch 

7 

8# Block variants with no attn_out + mlp_out split and no resid_mid hook. 

9_HYBRID_VARIANT_NAMES = ("mamba", "ssm", "mixer", "linear_attn") 

10_ID_DTYPES = (torch.int64, torch.int32, torch.int16, torch.int8, torch.uint8) 

11 

12 

13def require_eval_mode(model: Any, *, operation: str) -> None: 

14 """Reject training state anywhere in a wrapped model without mutating it.""" 

15 training_modules: dict[int, str] = {} 

16 roots = (("", model), ("original_model", getattr(model, "original_model", None))) 

17 for prefix, root in roots: 

18 if not isinstance(root, torch.nn.Module): 

19 continue 

20 for name, module in root.named_modules(): 

21 if not module.training: 

22 continue 

23 qualified_name = ".".join(part for part in (prefix, name) if part) 

24 training_modules.setdefault(id(module), qualified_name or "<root>") 

25 if not training_modules: 

26 return 

27 

28 names = list(training_modules.values()) 

29 preview = ", ".join(names[:3]) 

30 if len(names) > 3: 

31 preview += f", and {len(names) - 3} more" 

32 raise ValueError( 

33 f"{operation} requires the model and all submodules to be in evaluation " 

34 f"mode; found training mode at {preview}. Call model.eval() before running " 

35 "the analysis." 

36 ) 

37 

38 

39def hybrid_block_types(model: Any) -> List[str]: 

40 """Sorted distinct ``model.layer_types()`` entries naming an SSM or linear-attention block.""" 

41 layer_types = model.layer_types() if hasattr(model, "layer_types") else [] 

42 return sorted( 

43 {t for t in layer_types if any(part in _HYBRID_VARIANT_NAMES for part in t.split("+"))} 

44 ) 

45 

46 

47def compute_dtype(dtype: torch.dtype) -> torch.dtype: 

48 """Keep float64; run every other dtype in float32, which all backends' linalg support.""" 

49 return torch.float64 if dtype == torch.float64 else torch.float32 

50 

51 

52def require_int( 

53 value: Any, 

54 name: str, 

55 *, 

56 minimum: int = 1, 

57 maximum: Optional[int] = None, 

58 message: Optional[str] = None, 

59) -> int: 

60 """Return ``value`` as an int in ``[minimum, maximum]``; bools are refused as integers.""" 

61 if isinstance(value, bool) or not isinstance(value, Integral): 

62 raise ValueError(message or f"{name} must be an integer, got {value!r}") 

63 if value < minimum or (maximum is not None and value > maximum): 

64 if maximum is not None: 

65 expected = f"between {minimum} and {maximum}" 

66 elif minimum == 1: 

67 expected = "a positive integer" 

68 elif minimum == 0: 

69 expected = "a nonnegative integer" 

70 else: 

71 expected = f"an integer at least {minimum}" 

72 raise ValueError(message or f"{name} must be {expected}, got {value!r}") 

73 return int(value) 

74 

75 

76# Typed Any so the checks below, not a runtime type-hint checker, name a bad token. 

77def to_token_ids(model: Any, tokens: Any) -> List[int]: 

78 """Resolve one token, or a sequence of them, into in-vocabulary token ids. 

79 

80 A token is a single-token string or an integer id: a Python or numpy integer, or a 0-d 

81 integer tensor. A bare token and a one-element sequence holding it resolve alike. 

82 """ 

83 if isinstance(tokens, (str, Integral)) or ( 

84 isinstance(tokens, torch.Tensor) and tokens.ndim == 0 

85 ): 

86 tokens = [tokens] 

87 ids: List[int] = [] 

88 for token in tokens: 

89 if isinstance(token, torch.Tensor) and token.ndim == 0 and token.dtype in _ID_DTYPES: 

90 token = token.item() 

91 if isinstance(token, str): 

92 ids.append(int(model.to_single_token(token))) 

93 elif isinstance(token, bool): 

94 raise ValueError(f"token {token!r} must be a string or integer token id, not bool") 

95 elif isinstance(token, Integral): 

96 ids.append(int(token)) 

97 else: 

98 # int() would silently truncate a float id. 

99 raise TypeError( 

100 f"tokens must be strings or integer token ids, got {type(token).__name__}" 

101 ) 

102 if not ids: 

103 raise ValueError("tokens must contain at least one token") 

104 d_vocab = model.W_U.shape[1] 

105 invalid = [token_id for token_id in ids if not 0 <= token_id < d_vocab] 

106 if invalid: 

107 raise ValueError(f"token ids {invalid} out of range for vocabulary size {d_vocab}") 

108 return ids