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. It must be non-negative (a negative value raises ValueError), and it has no effect on greedy (temperature 0.0) sampling.

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 freq_penalty < 0:

101 # The penalty is applied only when positive, so a negative value would

102 # otherwise be a silent no-op rather than a frequency bonus.

103 raise ValueError(f"freq_penalty must be non-negative, got {freq_penalty}.")

104 if temperature == 0.0:

105 # Greedy sampling - still apply repetition penalty before argmax

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

107 final_logits = _apply_repetition_penalty(final_logits, tokens, repetition_penalty)

108 return final_logits.argmax(dim=-1)

109 else:

110 # Sample from the distribution

111

112 # Apply repetition penalty before temperature scaling

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

114 final_logits = _apply_repetition_penalty(final_logits, tokens, repetition_penalty)

115

116 final_logits = final_logits / temperature

117 if freq_penalty > 0:

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

119 assert (

120 len(tokens.shape) == 2

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

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

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

124 final_logits[batch_index] = final_logits[

125 batch_index

126 ] - freq_penalty * torch.bincount(

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

128 )

129 if top_k is not None:

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

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

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

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

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

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

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

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

138 elif top_p is not None:

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

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

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

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

143 sorted_indices_to_remove = cumulative_probs > top_p

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

145 sorted_indices_to_remove[..., 0] = 0

146 indices_to_remove = sorted_indices_to_remove.scatter(

147 -1, sorted_indices, sorted_indices_to_remove

148 )

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

150

151 final_logits = final_logits.to(torch.float32)

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

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

58 statements