Coverage for transformer_lens/utilities/lm_utils.py: 97%

1"""lm_utils. 

2 

3This module contains utility functions related to language models 

4""" 

5 

6from __future__ import annotations 

7 

8from typing import Optional, Union 

9 

10import torch 

11import torch.nn.functional as F 

12from jaxtyping import Bool, Float, Int 

13 

14 

15def lm_cross_entropy_loss( 

16 logits: Float[torch.Tensor, "batch pos d_vocab"], 

17 tokens: Int[torch.Tensor, "batch pos"], 

18 attention_mask: Optional[ 

19 Union[Bool[torch.Tensor, "batch pos"], Int[torch.Tensor, "batch pos"]] 

20 ] = None, 

21 per_token: bool = False, 

22) -> Union[Float[torch.Tensor, ""], Float[torch.Tensor, "batch pos-1"]]: 

23 """Cross entropy loss for the language model, gives the loss for predicting the NEXT token. 

24 

25 Args: 

26 logits (torch.Tensor): Logits. Shape [batch, pos, d_vocab] 

27 tokens (torch.Tensor[int64]): Input tokens. Shape [batch, pos] 

28 attention_mask (torch.Tensor[int64 or bool], optional): Attention mask. Shape [batch, pos]. 

29 Used to mask out padding tokens. Defaults to None. 

30 per_token (bool, optional): Whether to return the log probs predicted for the correct token, or the loss (ie mean of the predicted log probs). Note that the returned array has shape [batch, seq-1] as we cannot predict the first token (alternately, we ignore the final logit). Defaults to False. 

31 """ 

32 log_probs = F.log_softmax(logits, dim=-1) 

33 # Use torch.gather to find the log probs of the correct tokens 

34 # Offsets needed because we're predicting the NEXT token (this means the final logit is meaningless) 

35 # None and [..., 0] needed because the tensor used in gather must have the same rank. 

36 predicted_log_probs = log_probs[..., :-1, :].gather(dim=-1, index=tokens[..., 1:, None])[..., 0] 

37 

38 if attention_mask is not None: 

39 assert attention_mask.shape == tokens.shape, ( 

40 "attention_mask must have the same shape as tokens, " 

41 f"got {tuple(attention_mask.shape)} and {tuple(tokens.shape)}" 

42 ) 

43 # Ignore token positions which are masked out or where the next token is masked out 

44 # (generally padding tokens) 

45 next_token_mask = torch.logical_and(attention_mask[:, :-1], attention_mask[:, 1:]) 

46 predicted_log_probs = predicted_log_probs.masked_fill(~next_token_mask, 0.0) 

47 n_tokens = next_token_mask.sum().item() 

48 else: 

49 n_tokens = predicted_log_probs.numel() 

50 if per_token: 

51 return -predicted_log_probs 

52 else: 

53 return -predicted_log_probs.sum() / n_tokens 

54 

55 

56def lm_accuracy( 

57 logits: Float[torch.Tensor, "batch pos d_vocab"], 

58 tokens: Int[torch.Tensor, "batch pos"], 

59 per_token: bool = False, 

60) -> Union[Float[torch.Tensor, ""], Bool[torch.Tensor, "batch pos-1"]]: 

61 """Cross-Entropy Accuracy for Language Modelling. We measure the accuracy on the logits for predicting the NEXT token. 

62 

63 If per_token is True, returns the boolean for top 1 accuracy for each token in the batch. Note that this has size [batch, seq_len-1], as we cannot predict the first token. 

64 """ 

65 top_prediction = logits.argmax(dim=-1) 

66 correct_matches = top_prediction[:, :-1] == tokens[:, 1:] 

67 if per_token: 67 ↛ 70line 67 didn't jump to line 70 because the condition on line 67 was always true

68 return correct_matches 

69 else: 

70 return correct_matches.sum() / correct_matches.numel()