Coverage for transformer_lens/utilities/logits

1"""logits_utils.

3This module contains utility functions related to logits

4"""

6from __future__ import annotations

8from typing import Any, Optional

10import torch

11from jaxtyping import Float, Int

14def logits_to_df(

15 logits: Float[torch.Tensor, "d_vocab"],

16 tokenizer: Optional[Any] = None,

17 top_k: Optional[int] = None,

18) -> Any: # pandas.DataFrame; left as Any so beartype doesn't resolve a lazy import at runtime.

19 """Convert a 1-D logit vector into a sortable DataFrame for inspection.

21 Returns a frame with columns ``token_index``, ``token_string`` (when

22 ``tokenizer`` is given), ``logit``, ``log_prob``, ``probability``, sorted by

23 descending probability. ``top_k`` truncates to the highest-probability rows.

25 Args:

26 logits: 1-D tensor of shape [d_vocab]; raw model logits for one position.

27 tokenizer: Optional HF tokenizer used to materialise ``token_string``;

28 when ``None``, the column is omitted.

29 top_k: Optional cap on the number of returned rows.

30 """

31 # Lazy import — keeps `import transformer_lens` free of pandas's

32 # warnings unless logits_to_df is actually called.

33 import pandas as pd

35 log_probs = torch.log_softmax(logits.float(), dim=-1)

36 probs = log_probs.exp()

37 order = torch.argsort(probs, descending=True)

38 if top_k is not None:

39 order = order[:top_k]

41 indices = order.cpu().tolist()

42 data: dict = {"token_index": indices}

43 if tokenizer is not None:

44 data["token_string"] = [tokenizer.decode([i]) for i in indices]

45 data["logit"] = logits[order].detach().cpu().tolist()

46 data["log_prob"] = log_probs[order].detach().cpu().tolist()

47 data["probability"] = probs[order].detach().cpu().tolist()

48 return pd.DataFrame(data)

51def _apply_repetition_penalty(

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

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

54 penalty: float,

55) -> Float[torch.Tensor, "batch d_vocab"]:

56 """Apply HuggingFace-style repetition penalty to logits.

58 For each token that has appeared in the sequence, positive logits are divided

59 by the penalty and negative logits are multiplied by it.

61 Args:

62 logits: Logits tensor of shape [batch, d_vocab]

63 tokens: Token IDs of shape [batch, pos]

64 penalty: Repetition penalty value (1.0 = no penalty)

66 Returns:

67 Modified logits tensor

68 """

69 logits = logits.clone()

70 for batch_idx in range(logits.shape[0]):

71 unique_tokens = tokens[batch_idx].unique()

72 score = logits[batch_idx, unique_tokens]

73 # Divide positive logits, multiply negative logits

74 logits[batch_idx, unique_tokens] = torch.where(score > 0, score / penalty, score * penalty)

75 return logits

78def sample_logits(

79 final_logits: Float[torch.Tensor, "batch d_vocab"],

80 top_k: Optional[int] = None,

81 top_p: Optional[float] = None,

82 temperature: float = 1.0,

83 freq_penalty: float = 0.0,

84 repetition_penalty: float = 1.0,

85 tokens: Optional[Int[torch.Tensor, "batch pos"]] = None,

86) -> Int[torch.Tensor, "batch"]:

87 """

88 Sample from the logits, in order to generate text

90 final_logits has shape [batch, vocab_size]

91 We divide the logits by temperature before softmaxing and sampling - high temperature = more uniform, low = more argmaxy. Temp = 0.0 is greedy sampling

92 We apply top_k and top_p filtering to the logits, to encourage diversity. top_k = 10 means we only sample from the 10 most likely tokens. top_p = 0.9 means we only sample from the top 90% of tokens, and then renormalise the distribution. top_k and top_p are mutually exclusive. By default we apply neither and just sample from the full distribution.

94 Frequency penalty is a penalty on the probability of a token, proportional to the number of times it has been generated so far. This encourages the model to generate new tokens, rather than repeating itself. It is a hyperparameter, and should be tuned. It is applied to the logits before sampling. If this is non-zero it is required to input the input_tokens

96 Repetition penalty (HuggingFace-style) divides positive logits by the penalty value and multiplies negative logits by it for any token that has appeared in the sequence. A value of 1.0 means no penalty. Values > 1.0 discourage repetition. This is applied before temperature scaling.

98 When ``top_k`` exceeds the vocabulary size it is clamped to the vocabulary size (matching HuggingFace), rather than raising an error.

99 """

100 if temperature == 0.0:

101 # Greedy sampling - still apply repetition penalty before argmax

102 if repetition_penalty != 1.0 and tokens is not None:

103 final_logits = _apply_repetition_penalty(final_logits, tokens, repetition_penalty)

104 return final_logits.argmax(dim=-1)

105 else:

106 # Sample from the distribution

107

108 # Apply repetition penalty before temperature scaling

109 if repetition_penalty != 1.0 and tokens is not None: 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true

110 final_logits = _apply_repetition_penalty(final_logits, tokens, repetition_penalty)

111

112 final_logits = final_logits / temperature

113 if freq_penalty > 0:

114 assert tokens is not None, "Must provide input_tokens if applying a frequency penalty"

115 assert (

116 len(tokens.shape) == 2

117 ), "Frequency penalty do not support input in the form of embeddings"

118 for batch_index in range(final_logits.shape[0]):

119 # torch.bincount returns a tensor of length d_vocab, with the number of occurrences of each token in the tokens.

120 final_logits[batch_index] = final_logits[

121 batch_index

122 ] - freq_penalty * torch.bincount(

123 tokens[batch_index], minlength=final_logits.shape[-1]

124 )

125 if top_k is not None:

126 assert top_k > 0, "top_k has to be greater than 0"

127 # Clamp top_k to the vocab size so a large value does not raise

128 # "selected index k out of range" (matches HuggingFace's

129 # TopKLogitsWarper, which does top_k = min(top_k, logits.size(-1))).

130 top_k = min(top_k, final_logits.shape[-1])

131 top_logits, top_idx = final_logits.topk(top_k, dim=-1)

132 indices_to_remove = final_logits < top_logits[..., -1].unsqueeze(-1)

133 final_logits = final_logits.masked_fill(indices_to_remove, -float("inf"))

134 elif top_p is not None:

135 assert 1.0 >= top_p > 0.0, "top_p has to be in (0, 1]"

136 sorted_logits, sorted_indices = torch.sort(final_logits, descending=True)

137 cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

138 # We round up - we want prob >= top_p not <top_p

139 sorted_indices_to_remove = cumulative_probs > top_p

140 sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()

141 sorted_indices_to_remove[..., 0] = 0

142 indices_to_remove = sorted_indices_to_remove.scatter(

143 -1, sorted_indices, sorted_indices_to_remove

144 )

145 final_logits = final_logits.masked_fill(indices_to_remove, -float("inf"))

146

147 final_logits = final_logits.to(torch.float32)

148 return torch.distributions.categorical.Categorical(logits=final_logits).sample()

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

56 statements