stopping_criteria.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520
  1. import time
  2. import warnings
  3. from abc import ABC
  4. from collections import OrderedDict
  5. from copy import deepcopy
  6. from typing import Optional, Union
  7. import numpy as np
  8. import torch
  9. from torch.nn import functional as F
  10. from ..pytorch_utils import isin_mps_friendly
  11. from ..tokenization_utils_base import PreTrainedTokenizerBase
  12. from ..utils import add_start_docstrings, logging
  13. logger = logging.get_logger(__name__)
  14. # We maintain a module-level cache of the embedding vectors for the stop string criterion
  15. # because they are slow to compute
  16. STOP_STRING_EMBEDDING_CACHE = OrderedDict()
  17. STOPPING_CRITERIA_INPUTS_DOCSTRING = r"""
  18. Args:
  19. input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
  20. Indices of input sequence tokens in the vocabulary.
  21. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
  22. [`PreTrainedTokenizer.__call__`] for details.
  23. [What are input IDs?](../glossary#input-ids)
  24. scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
  25. Prediction scores of a language modeling head. These can be scores for each vocabulary token before SoftMax
  26. or scores for each vocabulary token after SoftMax. If this stopping criteria depends on the `scores` input,
  27. make sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`.
  28. kwargs (`dict[str, Any]`, *optional*):
  29. Additional stopping criteria specific kwargs.
  30. Return:
  31. `torch.BoolTensor`. (`torch.BoolTensor` of shape `(batch_size, 1)`):
  32. `True` indicates we stop generation for a particular row.
  33. `False` indicates we should continue.
  34. """
  35. class StoppingCriteria(ABC):
  36. """Abstract base class for all stopping criteria that can be applied during generation.
  37. If your stopping criteria depends on the `scores` input, make sure you pass `return_dict_in_generate=True,
  38. output_scores=True` to `generate`.
  39. """
  40. @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
  41. def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
  42. raise NotImplementedError("StoppingCriteria needs to be subclassed")
  43. class MaxLengthCriteria(StoppingCriteria):
  44. """
  45. This class can be used to stop generation whenever the full generated number of tokens exceeds `max_length`. Keep
  46. in mind for decoder-only type of transformers, this will include the initial prompted tokens.
  47. Args:
  48. max_length (`int`):
  49. The maximum length that the output sequence can have in number of tokens.
  50. max_position_embeddings (`int`, *optional*):
  51. The maximum model length, as defined by the model's `config.max_position_embeddings` attribute.
  52. """
  53. def __init__(self, max_length: int, max_position_embeddings: Optional[int] = None):
  54. self.max_length = max_length
  55. self.max_position_embeddings = max_position_embeddings
  56. @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
  57. def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
  58. cur_len = input_ids.shape[1]
  59. is_done = cur_len >= self.max_length
  60. if self.max_position_embeddings is not None and not is_done and cur_len > self.max_position_embeddings:
  61. logger.warning_once(
  62. "This is a friendly reminder - the current text generation call has exceeded the model's predefined "
  63. f"maximum length ({self.max_position_embeddings}). Depending on the model, you may observe "
  64. "exceptions, performance degradation, or nothing at all."
  65. )
  66. return torch.full((input_ids.shape[0],), is_done, device=input_ids.device, dtype=torch.bool)
  67. class MaxTimeCriteria(StoppingCriteria):
  68. """
  69. This class can be used to stop generation whenever the full generation exceeds some amount of time. By default, the
  70. time will start being counted when you initialize this function. You can override this by passing an
  71. `initial_time`.
  72. Args:
  73. max_time (`float`):
  74. The maximum allowed time in seconds for the generation.
  75. initial_time (`float`, *optional*, defaults to `time.time()`):
  76. The start of the generation allowed time.
  77. """
  78. def __init__(self, max_time: float, initial_timestamp: Optional[float] = None):
  79. self.max_time = max_time
  80. self.initial_timestamp = time.time() if initial_timestamp is None else initial_timestamp
  81. @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
  82. def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
  83. is_done = time.time() - self.initial_timestamp > self.max_time
  84. return torch.full((input_ids.shape[0],), is_done, device=input_ids.device, dtype=torch.bool)
  85. class StopStringCriteria(StoppingCriteria):
  86. """
  87. This class can be used to stop generation whenever specific string sequences are generated. It preprocesses
  88. the strings together with the tokenizer vocab to find positions where tokens can validly complete the stop strings.
  89. Generation is stopped as soon as a token is generated that completes any of the stop strings.
  90. We want to catch any instance in which the stop string would be present in the decoded output, which means
  91. we must also catch cases with "overhangs" off one or both ends. To make this more concrete, for the stop string
  92. "stop", any of the following token sequences would trigger the match:
  93. - ["st", "op"]
  94. - ["stop"]
  95. - ["st", "opera"]
  96. - ["sto", "pper"]
  97. - ["las", "topper"]
  98. - ["s", "to", "pped"]
  99. Note that a match will only be triggered if the stop string is at the end of the generated sequence. In other
  100. words, these sequences will not trigger a match:
  101. - ["stop", "at"]
  102. - ["st", "op", "at"]
  103. - ["st", "opera", "tion"]
  104. The reason these are not a match is that the stop string does not overlap with the final token. If you can remove
  105. one or more tokens from the end of the sequence without destroying the stop string, then this criterion will not
  106. match that stop string. This is by design; because this check is run after each token is generated, we can't miss a
  107. valid stop string if one is generated, but we don't want to halt generation just because the stop string exists
  108. somewhere in the past input_ids.
  109. How is the match actually performed, though? We do it in quite a confusing way, because we want the entire match
  110. process to be compilable with Torch or XLA, which means we cannot use standard string methods. However, it is possible,
  111. with some work, to do string matching with pure tensor operations. We'll begin by describing the algorithm we use
  112. with standard string operations, and then at the end we'll explain how this is converted to pure tensor operations.
  113. The key to the algorithm is an observation: Because the stop string must overlap with the end of the token sequence, we can start at
  114. the end of the sequence and work backwards. Specifically, we check that there is an overlap between the start of
  115. the final token and the end of the stop_string, or to put it another way, stop_string[-i:] == token[:i] for
  116. some i > 0. If you look at the positive examples above, you'll see the last token in all of them fulfills this
  117. property:
  118. - ["st", "op"] (overlap is "op", overlap length == 2)
  119. - ["stop"] (overlap is "stop", overlap length == 4)
  120. - ["st", "opera"] (overlap is "op", overlap length == 2)
  121. - ["sto", "pper"] (overlap is "p", overlap length == 1)
  122. - ["las", "topper"] (overlap is "top", overlap length == 3)
  123. - ["s", "to", "pped"] (overlap is "p", overlap length == 1)
  124. It's impossible to construct a matching sequence that does not have this property (feel free to verify this
  125. yourself). However, although this overlap between the start of the final token and the end of the stop string is
  126. necessary for a match, it is not sufficient. We also need to check that the rest of the token sequence is
  127. consistent with the stop string.
  128. How do we do that? Let's use ["s", "to", "pped"] as an example. We know that the final token, "pped", has an
  129. overlap of 1 with the stop string, "stop". We then go back to the previous token, "to". Since we have already
  130. matched 1 character from the stop string, the remainder to check is "sto". We check that the next token "to"
  131. matches the end of the remainder, which it does. We have now matched 3 characters from the stop string, and the
  132. remainder to match is "s". We go back to the previous token again, which is also "s". This is a match, and so
  133. we have matched the entire stop string.
  134. How does it work when the tokens run off the start of the stop string, though? Let's consider the example of
  135. ["las", "topper"]. The final token, "topper", has an overlap of 3 with the stop string, "stop". Therefore,
  136. the remaining stop string to match is "s". We go back to the previous token, "las". Because the remainder to
  137. match is just "s", with length 1, we consider only the final 1 character from the token, which is "s". This
  138. matches the stop string, and so the entire string is matched.
  139. How do we compute these matches with tensor operations, though? Simply: we efficiently precompute the necessary
  140. information for all tokens! For every token, we compute:
  141. - Its overlap with the end of the stop string, if any
  142. - The positions inside the stop string where the token matches, including matches that run off the start.
  143. - The total length of the token
  144. For example, for the token "pped", we would compute an end overlap of 1, no internal matching positions,
  145. and a length of 4. For the token "to", we would compute no end overlap, a single internal matching position
  146. of 1 (counting from the end), and a length of 2. For the token "s", we would compute no end overlap,
  147. a single internal matching position of 3 (again counting from the end) and a length of 1.
  148. As long as we have this information, we can execute the algorithm above without any string comparison
  149. operations. We simply perform the following steps:
  150. - Check if the final token has an end-overlap with the start string
  151. - Continue backwards, keeping track of how much of the stop string we've matched so far
  152. - At each point, check if the next token has the current position as one of its valid positions
  153. - Continue until either a match fails, or we completely match the whole stop string
  154. Again, consider ["s", "to", "pped"] as an example. "pped" has an end overlap of 1, so we can begin a match.
  155. We have matched 1 character so far, so we check that the next token "to", has 1 as a valid position (again,
  156. counting from the end). It does, so we add the length of "to" to our position tracker. We have now matched
  157. 3 characters, so we check that the next token "s" has 3 as a valid position. It does, so we add its length
  158. to the position tracker. The position tracker is now 4, which is the length of the stop string. We have matched the
  159. entire stop string.
  160. In the second case, ["las", "topper"], "topper" has an end overlap of 3, so we can begin a match. We have
  161. matched 3 characters so far, so we check that the next token "las" has 3 as a valid position. It does, because we
  162. allow tokens to match positions that run off the start of the stop string. We add its length to the position
  163. tracker. The position tracker is now 6, which is greater than the length of the stop string! Don't panic, though -
  164. this also counts as a match of the stop string. We have matched the entire stop string.
  165. Args:
  166. tokenizer (`PreTrainedTokenizer`):
  167. The model's associated tokenizer (necessary to extract vocab and tokenize the termination sequences)
  168. stop_strings (`Union[str, list[str]]`):
  169. A list of strings that should end generation. If a string is passed, it will be treated like a
  170. list with a single element.
  171. Examples:
  172. ```python
  173. >>> from transformers import AutoModelForCausalLM, AutoTokenizer
  174. >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
  175. >>> model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2")
  176. >>> inputs = tokenizer("The biggest states in the USA by land area:", return_tensors="pt")
  177. >>> gen_out = model.generate(**inputs)
  178. >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
  179. The biggest states in the USA by land area:
  180. - Alaska
  181. - Texas
  182. - California
  183. >>> # Passing one or more stop strings will halt generation after those strings are emitted
  184. >>> # Note that generating with stop strings requires you to pass the tokenizer too
  185. >>> gen_out = model.generate(**inputs, stop_strings=["Texas"], tokenizer=tokenizer)
  186. >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
  187. The biggest states in the USA by land area:
  188. - Alaska
  189. - Texas
  190. ```
  191. """
  192. def __init__(self, tokenizer: PreTrainedTokenizerBase, stop_strings: Union[str, list[str]]):
  193. if isinstance(stop_strings, str):
  194. stop_strings = [stop_strings]
  195. self.stop_strings: tuple[str, ...] = tuple(stop_strings)
  196. vocab = tokenizer.get_vocab()
  197. token_list, token_indices = tuple(vocab.keys()), tuple(vocab.values())
  198. self.embedding_vec, self.max_valid_positions, self.max_valid_end_lens = self.clean_and_embed_tokens_with_cache(
  199. token_list, token_indices, tokenizer
  200. )
  201. self.maximum_token_len = max(len(stop_string) for stop_string in self.stop_strings)
  202. self.num_stop_strings = len(self.stop_strings)
  203. self.target_lens = torch.tensor([len(stop_string) for stop_string in stop_strings], dtype=torch.int32)
  204. def clean_and_embed_tokens_with_cache(self, token_list, token_indices, tokenizer):
  205. # We don't use the tokenizer in the cache key, because I don't trust it to have well-behaved equality
  206. if (token_list, token_indices, self.stop_strings) in STOP_STRING_EMBEDDING_CACHE:
  207. embedding_vec, max_valid_positions, max_valid_end_lens = STOP_STRING_EMBEDDING_CACHE[
  208. (token_list, token_indices, self.stop_strings)
  209. ]
  210. STOP_STRING_EMBEDDING_CACHE.move_to_end((token_list, token_indices, self.stop_strings))
  211. else:
  212. clean_token_list, clean_token_indices = self.clean_tokenizer_vocab(tokenizer)
  213. embedding_vec, max_valid_positions, max_valid_end_lens = self._stop_string_create_embedding_vec(
  214. clean_token_list, clean_token_indices, self.stop_strings
  215. )
  216. STOP_STRING_EMBEDDING_CACHE[(token_list, token_indices, self.stop_strings)] = (
  217. embedding_vec,
  218. max_valid_positions,
  219. max_valid_end_lens,
  220. )
  221. if len(STOP_STRING_EMBEDDING_CACHE) > 8:
  222. STOP_STRING_EMBEDDING_CACHE.popitem(last=False) # Pop from the start, the least recently used item
  223. return embedding_vec, max_valid_positions, max_valid_end_lens
  224. @staticmethod
  225. def clean_tokenizer_vocab(tokenizer, static_prefix="abcdef"):
  226. """
  227. This method turns a tokenizer vocab into a "clean" vocab where each token represents the actual string
  228. it will yield, without any special prefixes like "##" or "Ġ". This is trickier than it looks - the method
  229. tokenizer.convert_tokens_to_string() does not always return the correct string because of issues with prefix
  230. space addition/removal. To work around this, we add a static prefix to the start of the token, then remove
  231. it (and any prefix that may have been introduced with it) after calling convert_tokens_to_string().
  232. """
  233. vocab = tokenizer.get_vocab()
  234. clean_token_list = []
  235. clean_token_indices = []
  236. sentence_base = tokenizer(static_prefix, add_special_tokens=False)["input_ids"]
  237. tokens_base = [tokenizer._convert_id_to_token(tok) for tok in sentence_base]
  238. for token, token_idx in vocab.items():
  239. token_string = tokenizer.convert_tokens_to_string(tokens_base + [token])
  240. token_string = token_string[token_string.index(static_prefix) + len(static_prefix) :]
  241. clean_token_list.append(token_string)
  242. clean_token_indices.append(token_idx)
  243. return tuple(clean_token_list), tuple(clean_token_indices)
  244. @staticmethod
  245. def _stop_string_get_matching_positions(
  246. token_list, token_indices, stop_strings
  247. ) -> tuple[dict[str, dict[str, list[int]]], dict[str, dict[str, list[int]]]]:
  248. """This function preprocesses stop strings and the tokenizer vocabulary to determine where tokens can
  249. validly appear in the stop strings. For each token, it computes a list of positions in the stop string where the
  250. token appears, as well as a list of the possible "end overlaps" for that token - that is, the number of characters
  251. from the end of the stop string that overlap with the start of the token, which can have more than one value.
  252. The reason for computing these may seem a bit cryptic - please see the docstring for StopStringCriteria for a full
  253. explanation of what these values are for!"""
  254. token_valid_positions = {}
  255. token_end_overlaps = {}
  256. for stop_string in stop_strings:
  257. reversed_stop_string = stop_string[::-1]
  258. token_valid_positions[stop_string] = {}
  259. token_end_overlaps[stop_string] = {}
  260. for token, tok_idx in zip(token_list, token_indices):
  261. reversed_token = token[::-1]
  262. matching_positions = []
  263. possible_end_lengths = []
  264. for i in range(1 - len(token), len(stop_string)):
  265. if i < 0:
  266. tok = reversed_token[-i:]
  267. i = 0
  268. else:
  269. tok = reversed_token
  270. stop = reversed_stop_string[i : i + len(tok)]
  271. if tok.startswith(stop):
  272. if i == 0:
  273. possible_end_lengths.append(min(len(tok), len(stop)))
  274. else:
  275. matching_positions.append(i)
  276. if matching_positions:
  277. token_valid_positions[stop_string][tok_idx] = matching_positions
  278. if possible_end_lengths:
  279. token_end_overlaps[stop_string][tok_idx] = possible_end_lengths
  280. return token_valid_positions, token_end_overlaps
  281. @staticmethod
  282. def _stop_string_create_embedding_vec(token_list, token_indices, stop_strings) -> dict[str, torch.tensor]:
  283. """This function precomputes everything needed for the run-time checks in StopStringCriteria, and packs
  284. them into an embedding tensor that can be accessed with pure tensor operations. For the specifics of the values
  285. that are precomputed and what they are used for, please refer to the StopStringCriteria docstring!"""
  286. token_valid_positions, token_end_overlaps = StopStringCriteria._stop_string_get_matching_positions(
  287. token_list, token_indices, stop_strings
  288. )
  289. all_valid_positions = [len(val) for positions in token_valid_positions.values() for val in positions.values()]
  290. # In some cases, tokens may have no valid internal positions (such as single-character stop strings), so
  291. # we need a fallback to handle this case
  292. max_valid_positions = max(all_valid_positions) if all_valid_positions else 1
  293. # There should always be at least one valid end_len, however, so no fallback needed here
  294. valid_end_lens = [len(val) for positions in token_end_overlaps.values() for val in positions.values()]
  295. if not valid_end_lens:
  296. raise ValueError(
  297. "Stop string preprocessing was unable to identify tokens matching one or more of the "
  298. "supplied stop string(s). This is most often caused by the stop "
  299. "strings containing unusual characters that are not in the tokenizer vocabulary."
  300. )
  301. max_valid_end_lens = max(valid_end_lens)
  302. vec_size = len(stop_strings) * (max_valid_positions + max_valid_end_lens) + 1
  303. # We use +2 instead of +1 so we can have a dummy entry at the end. We will clamp all token values
  304. # over the max to this, ensuring they do not contribute to stop string matching.
  305. gather_vec = np.full((max(token_indices) + 2, vec_size), dtype=np.int32, fill_value=-1)
  306. for i, stop_string in enumerate(stop_strings):
  307. positions = token_valid_positions[stop_string]
  308. end_lens = token_end_overlaps[stop_string]
  309. # Since this is lots of very small assignments of lists, we build it with numpy rather
  310. # than torch for speed + simplicity, then convert to torch at the end
  311. for token_idx, valid_positions in positions.items():
  312. gather_vec[token_idx, max_valid_positions * i : max_valid_positions * i + len(valid_positions)] = (
  313. valid_positions
  314. )
  315. for token_idx, possible_end_lens in end_lens.items():
  316. gather_vec[
  317. token_idx,
  318. max_valid_positions * len(stop_strings) + max_valid_end_lens * i : max_valid_positions
  319. * len(stop_strings)
  320. + max_valid_end_lens * i
  321. + len(possible_end_lens),
  322. ] = possible_end_lens
  323. for token, token_idx in zip(token_list, token_indices):
  324. gather_vec[token_idx, -1] = len(token)
  325. gather_vec = torch.tensor(gather_vec, dtype=torch.int32)
  326. return gather_vec, max_valid_positions, max_valid_end_lens
  327. @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
  328. def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.Tensor:
  329. self.embedding_vec = self.embedding_vec.to(input_ids.device)
  330. self.target_lens = self.target_lens.to(input_ids.device)
  331. # The maximum length we need to consider is 1 token per character. Note that input_ids can also be
  332. # *shorter* than the global max, and the code below should be ready for that
  333. input_ids = input_ids[:, -self.maximum_token_len :]
  334. # Flip input_ids because we're only matching strings at the end of the generated sequence
  335. flipped_ids = torch.flip(input_ids, (1,))
  336. # Clip out-of-vocab values to the dummy value at the end of the embedding vector
  337. flipped_ids = torch.clamp(flipped_ids, max=self.embedding_vec.size(0) - 1)
  338. # Size of the vector of positions a single token can match
  339. max_valid_positions = self.max_valid_positions
  340. # The embedding vec contains the valid positions, end_lengths and total lengths for each token
  341. embedded = F.embedding(flipped_ids, self.embedding_vec)
  342. # Now we split the embedding vector. valid_positions is the positions in the stop string the token can fit
  343. valid_positions = embedded[:, 1:, : max_valid_positions * self.num_stop_strings].unflatten(
  344. -1, (self.num_stop_strings, -1)
  345. )
  346. # end_lengths is the number of characters from the string, counting from the end, that the token
  347. # contains. It can have multiple values if the same token can overlap different end lengths
  348. end_lengths = embedded[:, :1, max_valid_positions * self.num_stop_strings : -1].unflatten(
  349. -1, (self.num_stop_strings, -1)
  350. )
  351. # Lengths is the total length of each token. Unlike the others, it always has a single value
  352. lengths = embedded[:, 1:, None, -1:] # Insert a dummy dimension for stop_strings even though lengths are const
  353. # Concatenate lengths onto each possible end_lengths value
  354. lengths = lengths.expand((-1, -1, end_lengths.shape[-2], end_lengths.shape[-1]))
  355. lengths_with_ends = torch.cat([end_lengths, lengths], dim=1)
  356. # cumsum() to get the number of matched characters in the stop string after each token
  357. cumsum = lengths_with_ends.cumsum(dim=1) # B x maximum_token_len x num_stop_strings x max_valid_end_lens
  358. # The calculation above assumes that all tokens are in valid positions. Now we mask the ones that are not.
  359. # First, tokens match the start of the string if they have a positive value in the end_lengths vector
  360. initial_match = end_lengths > 0
  361. # Tokens continue the string if the cumsum() so far is one of the valid positions for that token
  362. # Note that we're actually tracking one cumsum() for for each possible end_length
  363. later_match = torch.any(cumsum[:, :-1, :, None] == valid_positions[:, :, :, :, None], axis=-2)
  364. # The match vector is a boolean vector that indicates which positions have valid tokens
  365. match = torch.cat([initial_match, later_match], dim=1)
  366. # Once a single position does not match, all positions following that position are masked
  367. mask = (~match).cumsum(dim=1, dtype=torch.int32)
  368. mask = mask == 0
  369. # The string is matched if we reached a cumsum equal to or greater than the length of the string
  370. # before hitting the mask
  371. string_matches = torch.amax(cumsum * mask, dim=(1, -1)) >= self.target_lens[None, :]
  372. # We return a per-sample vector that is True if any stop string is matched for that sample
  373. return torch.any(string_matches, dim=-1)
  374. class EosTokenCriteria(StoppingCriteria):
  375. """
  376. This class can be used to stop generation whenever the "end-of-sequence" token is generated.
  377. By default, it uses the `model.generation_config.eos_token_id`.
  378. Args:
  379. eos_token_id (`Union[int, list[int], torch.Tensor]`):
  380. The id(s) of the *end-of-sequence* token.
  381. """
  382. def __init__(self, eos_token_id: Union[int, list[int], torch.Tensor]):
  383. if not isinstance(eos_token_id, torch.Tensor):
  384. if isinstance(eos_token_id, int):
  385. eos_token_id = [eos_token_id]
  386. eos_token_id = torch.tensor(eos_token_id)
  387. self.eos_token_id = eos_token_id
  388. @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
  389. def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
  390. self.eos_token_id = self.eos_token_id.to(input_ids.device)
  391. is_done = isin_mps_friendly(input_ids[:, -1], self.eos_token_id)
  392. return is_done
  393. class ConfidenceCriteria(StoppingCriteria):
  394. """
  395. This class can be used to stop generation whenever assistant model's confidence in its prediction for the current token is lower than the threshold
  396. `model.generation_config.assistant_confidence_threshold` even if the number of speculative tokens (defined by `num_assistant_tokens`) is not yet reached.
  397. Args:
  398. assistant_confidence_threshold (`float`):
  399. The value of the threshold.
  400. """
  401. def __init__(self, assistant_confidence_threshold):
  402. self.assistant_confidence_threshold = assistant_confidence_threshold
  403. def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
  404. probs = scores[-1].softmax(-1)
  405. p = probs[0, input_ids[0, -1]].item()
  406. if p < self.assistant_confidence_threshold:
  407. return True
  408. return False
  409. class StoppingCriteriaList(list):
  410. @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
  411. def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
  412. is_done = torch.full((input_ids.shape[0],), False, device=input_ids.device, dtype=torch.bool)
  413. for criteria in self:
  414. is_done = is_done | criteria(input_ids, scores, **kwargs)
  415. return is_done
  416. @property
  417. def max_length(self) -> Optional[int]:
  418. for stopping_criterium in self:
  419. if isinstance(stopping_criterium, MaxLengthCriteria):
  420. return stopping_criterium.max_length
  421. return None
  422. def validate_stopping_criteria(stopping_criteria: StoppingCriteriaList, max_length: int) -> StoppingCriteriaList:
  423. stopping_max_length = stopping_criteria.max_length
  424. new_stopping_criteria = deepcopy(stopping_criteria)
  425. if stopping_max_length is not None and stopping_max_length != max_length:
  426. warnings.warn("You set different `max_length` for stopping criteria and `max_length` parameter", UserWarning)
  427. elif stopping_max_length is None:
  428. new_stopping_criteria.append(MaxLengthCriteria(max_length=max_length))
  429. return new_stopping_criteria