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

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

2 

3from typing import Any 

4 

5import torch 

6 

7 

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

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

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

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

12 for prefix, root in roots: 

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

14 continue 

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

16 if not module.training: 

17 continue 

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

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

20 if not training_modules: 

21 return 

22 

23 names = list(training_modules.values()) 

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

25 if len(names) > 3: 

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

27 raise ValueError( 

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

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

30 "the analysis." 

31 )