Coverage for transformer_lens/model_bridge/sources/inspect/masks.py: 100%

1"""Torch-free validation of the single-sequence Inspect padding-mask wire format.""" 

2 

3from typing import Any 

4 

5import numpy as np 

6 

7 

8def normalize_attention_mask(mask: Any, n_tokens: int) -> list[int]: 

9 """Accept a binary flat wire mask or a single-row public padding mask.""" 

10 if hasattr(mask, "detach"): 

11 mask = mask.detach().cpu().tolist() 

12 array = np.asarray(mask) 

13 if array.ndim == 2 and array.shape[0] == 1: 

14 array = array[0] 

15 if array.ndim != 1 or array.shape[0] != n_tokens: 

16 raise ValueError("Inspect attention_mask must match the single input sequence length.") 

17 if not np.all((array == 0) | (array == 1)): 

18 raise ValueError("Inspect attention_mask must be a binary 0/1 padding mask.") 

19 if not np.any(array): 

20 raise ValueError("Inspect attention_mask must retain at least one token.") 

21 return [int(value) for value in array.tolist()]