Coverage for transformer_lens/lit/utils.py: 55%
122 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Utility functions for the LIT integration module.
3This module provides helper functions for converting between TransformerLens
4data structures and LIT-compatible formats, as well as other utilities.
6References:
7 - LIT API: https://pair-code.github.io/lit/documentation/api
8 - TransformerLens: https://github.com/TransformerLensOrg/TransformerLens
9"""
11from __future__ import annotations
13import logging
14from typing import Any, Dict, List, Optional, Tuple, Union
16import numpy as np
17import torch
19logger = logging.getLogger(__name__)
22def check_lit_installed() -> bool:
23 """Check if LIT (lit-nlp) is installed.
25 Returns:
26 bool: True if LIT is installed, False otherwise.
27 """
28 try:
29 import lit_nlp # noqa: F401
31 return True
32 except ImportError:
33 return False
36def tensor_to_numpy(
37 tensor: Union[torch.Tensor, np.ndarray, None],
38) -> Optional[np.ndarray]:
39 """Convert a PyTorch tensor to a NumPy array.
41 LIT expects all data to be in NumPy format, so this helper ensures
42 proper conversion with detach and CPU transfer.
44 Args:
45 tensor: PyTorch tensor or None.
47 Returns:
48 NumPy array or None if input was None.
49 """
50 if tensor is None:
51 return None
52 if isinstance(tensor, np.ndarray):
53 return tensor
54 if isinstance(tensor, torch.Tensor): 54 ↛ 56line 54 didn't jump to line 56 because the condition on line 54 was always true
55 return tensor.detach().cpu().numpy()
56 raise TypeError(f"Expected torch.Tensor or np.ndarray, got {type(tensor)}")
59def numpy_to_tensor(
60 array: Union[np.ndarray, torch.Tensor, None],
61 device: Optional[Union[str, torch.device]] = None,
62 dtype: Optional[torch.dtype] = None,
63) -> Optional[torch.Tensor]:
64 """Convert a NumPy array to a PyTorch tensor.
66 Args:
67 array: NumPy array or None.
68 device: Target device for the tensor.
69 dtype: Target dtype for the tensor.
71 Returns:
72 PyTorch tensor or None if input was None.
73 """
74 if array is None: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true
75 return None
76 if isinstance(array, torch.Tensor): 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true
77 tensor = array
78 else:
79 tensor = torch.from_numpy(array)
81 if dtype is not None: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true
82 tensor = tensor.to(dtype)
83 if device is not None:
84 tensor = tensor.to(device)
85 return tensor
88def get_tokens_from_model(
89 model: Any,
90 text: str,
91 prepend_bos: bool = True,
92 truncate: bool = True,
93 max_length: Optional[int] = None,
94) -> Tuple[List[str], torch.Tensor]:
95 """Get tokens and token IDs from a HookedTransformer model.
97 Args:
98 model: HookedTransformer model with tokenizer.
99 text: Input text to tokenize.
100 prepend_bos: Whether to prepend the BOS token.
101 truncate: Whether to truncate to max_length.
102 max_length: Maximum sequence length.
104 Returns:
105 Tuple of (token strings, token ID tensor).
107 Raises:
108 ValueError: If model has no tokenizer.
109 """
110 if model.tokenizer is None:
111 raise ValueError("Model must have a tokenizer to convert text to tokens")
113 token_ids = model.to_tokens(text, prepend_bos=prepend_bos, truncate=truncate)
115 if max_length is not None and token_ids.shape[1] > max_length:
116 token_ids = token_ids[:, :max_length]
118 token_strings = model.tokenizer.convert_ids_to_tokens(token_ids.squeeze(0).tolist())
120 return token_strings, token_ids.squeeze(0)
123def clean_token_string(token: str) -> str:
124 """Clean a token string for display.
126 Handles common tokenizer artifacts like:
127 - Ġ (GPT-2 style space prefix)
128 - ▁ (SentencePiece space prefix)
129 - ## (BERT style subword prefix)
131 Args:
132 token: Raw token string from tokenizer.
134 Returns:
135 Cleaned token string for display.
136 """
137 # Handle GPT-2/RoBERTa style space encoding
138 if token.startswith("Ġ"):
139 return "▁" + token[1:] # Use Unicode space indicator
140 # Handle SentencePiece
141 if token.startswith("▁"):
142 return token # Already in preferred format
143 # Handle BERT style
144 if token.startswith("##"):
145 return token[2:] # Remove ## prefix
146 return token
149def clean_token_strings(tokens: List[str]) -> List[str]:
150 """Clean a list of token strings for display.
152 Args:
153 tokens: List of raw token strings.
155 Returns:
156 List of cleaned token strings.
157 """
158 return [clean_token_string(t) for t in tokens]
161def extract_attention_from_cache(
162 cache: Any,
163 layer: int,
164 head: Optional[int] = None,
165 batch_idx: int = 0,
166) -> Optional[np.ndarray]:
167 """Extract attention patterns from an activation cache.
169 Args:
170 cache: TransformerLens ActivationCache object.
171 layer: Layer index to extract from.
172 head: Optional head index. If None, returns all heads.
173 batch_idx: Batch index to extract.
175 Returns:
176 Attention pattern as numpy array.
177 Shape: [query_pos, key_pos] if head specified
178 Shape: [num_heads, query_pos, key_pos] if head is None
179 """
180 attn_pattern = cache[f"blocks.{layer}.attn.hook_pattern"]
182 # Remove batch dimension
183 if attn_pattern.dim() == 4:
184 attn_pattern = attn_pattern[batch_idx]
186 # attn_pattern shape: [num_heads, query_pos, key_pos]
187 if head is not None:
188 attn_pattern = attn_pattern[head]
190 return tensor_to_numpy(attn_pattern)
193def extract_embeddings_from_cache(
194 cache: Any,
195 layer: int,
196 position: str = "all",
197 batch_idx: int = 0,
198) -> Optional[np.ndarray]:
199 """Extract embeddings from a specific layer in the activation cache.
201 Args:
202 cache: TransformerLens ActivationCache object.
203 layer: Layer index to extract from.
204 position: "all" for all positions, "first" for CLS-like, "last" for final token.
205 batch_idx: Batch index to extract.
207 Returns:
208 Embeddings as numpy array.
209 """
210 resid = cache[f"blocks.{layer}.hook_resid_post"]
212 # Remove batch dimension
213 if resid.dim() == 3:
214 resid = resid[batch_idx]
216 # resid shape: [seq_len, d_model]
217 if position == "first":
218 embeddings = resid[0]
219 elif position == "last":
220 embeddings = resid[-1]
221 elif position == "mean":
222 embeddings = resid.mean(dim=0)
223 else: # "all"
224 embeddings = resid
226 return tensor_to_numpy(embeddings)
229def compute_token_gradients(
230 model: Any,
231 text: str,
232 target_idx: Optional[int] = None,
233 prepend_bos: bool = True,
234) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List[str]]:
235 """Compute token-level gradients for salience.
237 Uses gradient of the loss with respect to token embeddings to compute
238 importance scores for each token.
240 Args:
241 model: HookedTransformer model.
242 text: Input text.
243 target_idx: Target token index for gradient computation.
244 If None, uses the last token.
245 prepend_bos: Whether to prepend BOS token.
247 Returns:
248 Tuple of (grad_l2, grad_dot_input, tokens) where:
249 - grad_l2: L2 norm of gradients per token [seq_len]
250 - grad_dot_input: Gradient dot input embedding per token [seq_len]
251 - tokens: List of token strings
252 """
253 # Tokenize
254 tokens, token_ids = get_tokens_from_model(model, text, prepend_bos=prepend_bos)
255 token_ids = token_ids.unsqueeze(0).to(model.cfg.device)
257 input_embeds = model.embed(token_ids)
258 input_embeds.requires_grad_(True)
260 # Forward pass
261 logits = model(input_embeds, start_at_layer=0)
263 # Determine target
264 if target_idx is None:
265 target_idx = -1 # Last token
267 # Get target logit and compute gradient
268 target_logit = logits[0, target_idx, token_ids[0, target_idx + 1]]
269 target_logit.backward()
271 gradients = input_embeds.grad[0] # [seq_len, d_model]
273 # Compute gradient L2 norm per token
274 grad_l2 = torch.norm(gradients, dim=-1) # [seq_len]
276 # Compute gradient dot input
277 grad_dot_input = (gradients * input_embeds[0].detach()).sum(dim=-1) # [seq_len]
279 return (
280 tensor_to_numpy(grad_l2),
281 tensor_to_numpy(grad_dot_input),
282 tokens,
283 )
286def get_top_k_predictions(
287 logits: torch.Tensor,
288 tokenizer: Any,
289 k: int = 10,
290 position: int = -1,
291 batch_idx: int = 0,
292) -> List[Tuple[str, float]]:
293 """Get top-k token predictions with their probabilities.
295 Args:
296 logits: Model logits tensor.
297 tokenizer: HuggingFace tokenizer.
298 k: Number of top predictions to return.
299 position: Position index to get predictions for.
300 batch_idx: Batch index.
302 Returns:
303 List of (token_string, probability) tuples.
304 """
305 pos_logits = logits[batch_idx, position] # [d_vocab]
307 probs = torch.softmax(pos_logits, dim=-1)
309 top_probs, top_indices = torch.topk(probs, k)
311 # Convert to strings
312 results = []
313 for prob, idx in zip(top_probs.tolist(), top_indices.tolist()):
314 token_str = tokenizer.decode([idx])
315 results.append((token_str, prob))
317 return results
320def validate_input_example(
321 example: Dict[str, Any],
322 required_fields: List[str],
323) -> bool:
324 """Validate that an input example has all required fields.
326 Args:
327 example: Input example dictionary.
328 required_fields: List of required field names.
330 Returns:
331 True if valid, False otherwise.
332 """
333 for field in required_fields:
334 if field not in example:
335 logger.warning(f"Missing required field '{field}' in input example")
336 return False
337 return True
340def batch_examples(
341 examples: List[Dict[str, Any]],
342 batch_size: int,
343) -> List[List[Dict[str, Any]]]:
344 """Split examples into batches.
346 Args:
347 examples: List of example dictionaries.
348 batch_size: Size of each batch.
350 Returns:
351 List of batches, where each batch is a list of examples.
352 """
353 return [examples[i : i + batch_size] for i in range(0, len(examples), batch_size)]
356def unbatch_outputs(
357 batched_outputs: Dict[str, np.ndarray],
358) -> List[Dict[str, Any]]:
359 """Split batched outputs into individual examples.
361 Takes a dictionary with batched arrays and returns a list of
362 dictionaries with individual arrays.
364 Args:
365 batched_outputs: Dictionary mapping field names to batched arrays.
367 Returns:
368 List of dictionaries, one per example.
369 """
370 if not batched_outputs: 370 ↛ 371line 370 didn't jump to line 371 because the condition on line 370 was never true
371 return []
373 # Get batch size from first array
374 first_key = next(iter(batched_outputs))
375 batch_size = len(batched_outputs[first_key])
377 # Split into individual examples
378 results = []
379 for i in range(batch_size):
380 example_output = {}
381 for key, value in batched_outputs.items():
382 if isinstance(value, (np.ndarray, torch.Tensor)):
383 example_output[key] = value[i]
384 elif isinstance(value, list): 384 ↛ 387line 384 didn't jump to line 387 because the condition on line 384 was always true
385 example_output[key] = value[i]
386 else:
387 example_output[key] = value
388 results.append(example_output)
390 return results
393def get_hook_name_for_layer(template: str, layer: int, **kwargs) -> str:
394 """Generate a hook point name from a template.
396 Args:
397 template: Hook name template with {layer} placeholder.
398 layer: Layer index.
399 **kwargs: Additional template parameters.
401 Returns:
402 Formatted hook point name.
403 """
404 return template.format(layer=layer, **kwargs)
407def filter_cache_by_pattern(
408 cache: Any,
409 pattern: str,
410) -> Dict[str, torch.Tensor]:
411 """Filter activation cache entries by hook name pattern.
413 Args:
414 cache: TransformerLens ActivationCache.
415 pattern: Pattern to match (e.g., "attn.hook_pattern" will match
416 all attention pattern hooks).
418 Returns:
419 Dictionary of matching cache entries.
420 """
421 return {name: value for name, value in cache.items() if pattern in name}
424def get_model_info(model: Any) -> Dict[str, Any]:
425 """Extract relevant model information for LIT display.
427 Args:
428 model: HookedTransformer model.
430 Returns:
431 Dictionary with model metadata.
432 """
433 cfg = model.cfg
434 return {
435 "model_name": cfg.model_name,
436 "n_layers": cfg.n_layers,
437 "n_heads": cfg.n_heads,
438 "d_model": cfg.d_model,
439 "d_head": cfg.d_head,
440 "d_mlp": cfg.d_mlp,
441 "d_vocab": cfg.d_vocab,
442 "n_ctx": cfg.n_ctx,
443 "act_fn": cfg.act_fn,
444 "normalization_type": cfg.normalization_type,
445 "positional_embedding_type": cfg.positional_embedding_type,
446 }