sampling_params.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456
  1. """Sampling parameters for text generation."""
  2. import copy
  3. import os
  4. from enum import IntEnum
  5. from functools import cached_property
  6. from typing import Any, Callable, Dict, List, Optional, Union
  7. import torch
  8. from loguru import logger
  9. from pydantic import Field
  10. from typing_extensions import Annotated
  11. _SAMPLING_EPS = 1e-5
  12. APHRODITE_NO_DEPRECATION_WARNING = bool(
  13. int(os.environ.get("APHRODITE_NO_DEPRECATION_WARNING", "0")))
  14. class SamplingType(IntEnum):
  15. GREEDY = 0
  16. RANDOM = 1
  17. RANDOM_SEED = 2
  18. BEAM = 3
  19. LogitsProcessorFunc = Union[Callable[[List[int], torch.Tensor], torch.Tensor],
  20. Callable[[List[int], List[int], torch.Tensor],
  21. torch.Tensor]]
  22. """LogitsProcessor is a function that takes a list
  23. of previously generated tokens, the logits tensor
  24. for the next token and, optionally, prompt tokens as a
  25. first argument, and returns a modified tensor of logits
  26. to sample from."""
  27. class SamplingParams:
  28. """Sampling parameters for text generation.
  29. Overall, we follow the sampling parameters from the OpenAI text completion
  30. API (https://platform.openai.com/docs/api-reference/completions/create).
  31. In addition, we support multiple additional samplers which are not supported
  32. by OpenAI.
  33. Args:
  34. n: Number of output sequences to return for the given prompt.
  35. best_of: Number of output sequences that are generated from the prompt.
  36. From these `best_of` sequences, the top `n` sequences are returned.
  37. `best_of` must be greater than or equal to `n`. This is treated as
  38. the beam width when `use_beam_search` is True. By default, `best_of`
  39. is set to `n`.
  40. presence_penalty: Float that penalizes new tokens based on whether they
  41. appear in the generated text so far. Values > 0 encourage the model
  42. to use new tokens, while values < 0 encourage the model to repeat
  43. tokens.
  44. frequency_penalty: Float that penalizes new tokens based on their
  45. frequency in the generated text so far. Values > 0 encourage the
  46. model to use new tokens, while values < 0 encourage the model to
  47. repeat tokens.
  48. repetition_penalty: Float that penalizes new tokens based on their
  49. frequency in the generated text so far.
  50. freq_pen is applied additively while
  51. rep_pen is applied multiplicatively.
  52. Must be in [1, inf). Set to 1 to disable the effect.
  53. temperature: Float that controls the randomness of the sampling. Lower
  54. values make the model more deterministic, while higher values make
  55. the model more random. Zero means greedy sampling.
  56. top_p: Float that controls the cumulative probability of the top tokens
  57. to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
  58. top_k: Integer that controls the number of top tokens to consider. Set
  59. to -1 to consider all tokens.
  60. top_a: Float that controls the cutoff for Top-A sampling.
  61. Exact cutoff is top_a*max_prob**2. Must be in [0,inf], 0 to disable.
  62. min_p: Float that controls the cutoff for min-p sampling.
  63. Exact cutoff is min_p*max_prob. Must be in [0,1], 0 to disable.
  64. tfs: Float that controls the cumulative approximate curvature of the
  65. distribution to retain for Tail Free Sampling.
  66. Must be in (0, 1]. Set to 1 to disable
  67. eta_cutoff: Float that controls the cutoff threshold for Eta sampling
  68. (a form of entropy adaptive truncation sampling)
  69. threshold is computed as min(eta, sqrt(eta)*entropy(probs)).
  70. Specified in units of 1e-4. Set to 0 to disable
  71. epsilon_cutoff: Float that controls the cutoff threshold for
  72. Epsilon sampling (simple probability threshold truncation).
  73. Specified in units of 1e-4. Set to 0 to disable.
  74. typical_p: Float that controls the cumulative probability of tokens
  75. closest in surprise to the expected surprise to consider.
  76. Must be in (0, 1]. Set to 1 to disable.
  77. mirostat_mode: Can either be 0 (disabled) or 2 (Mirostat v2).
  78. mirostat_tau: Target "surprisal" that mirostat works towards.
  79. Range [0, inf).
  80. mirostat_eta: Rate at which mirostat updates its internal surprisal
  81. value. Range [0, inf).
  82. dynatemp_min: Minimum temperature for dynatemp sampling.
  83. Range [0, inf).
  84. dynatemp_max: Maximum temperature for dynatemp sampling.
  85. Range [0, inf).
  86. dynatemp_exponent: Exponent for dynatemp sampling. Range [0, inf).
  87. smoothing_factor: Smoothing factor for Quadratic Sampling.
  88. smoothing_curve: Smoothing curve for Quadratic (Cubic) Sampling.
  89. seed: Random seed to use for the generation.
  90. use_beam_search: Whether to use beam search instead of sampling.
  91. length_penalty: Float that penalizes sequences based on their length.
  92. Used in beam search.
  93. early_stopping: Controls the stopping condition for beam search. It
  94. accepts the following values: `True`, where the generation stops as
  95. soon as there are `best_of` complete candidates; `False`, where an
  96. heuristic is applied and the generation stops when is it very
  97. unlikely to find better candidates; `"never"`, where the beam search
  98. procedure only stops when there cannot be better candidates
  99. (canonical beam search algorithm).
  100. stop: List of strings that stop the generation when they are generated.
  101. The returned output will not contain the stop strings.
  102. stop_token_ids: List of tokens that stop the generation when they are
  103. generated. The returned output will contain the stop tokens unless
  104. the stop tokens are special tokens.
  105. include_stop_str_in_output: Whether to include the stop strings in
  106. output text. Defaults to False.
  107. ignore_eos: Whether to ignore the EOS token and continue generating
  108. tokens after the EOS token is generated.
  109. max_tokens: Maximum number of tokens to generate per output sequence.
  110. min_tokens: Minimum number of tokens to generate per output sequence
  111. before EOS or stop tokens are generated.
  112. logprobs: Number of log probabilities to return per output token.
  113. When set to None, no probability is returned. If set to a non-None
  114. value, the result includes the log probabilities of the specified
  115. number of most likely tokens, as well as the chosen tokens.
  116. Note that the implementation follows the OpenAI API: The API will
  117. always return the log probability of the sampled token, so there
  118. may be up to `logprobs+1` elements in the response.
  119. prompt_logprobs: Number of log probabilities to return per prompt token.
  120. detokenize: Whether to detokenize the output. Defaults to True.
  121. custom_token_bans: List of token IDs to ban from generating
  122. skip_special_tokens: Whether to skip special tokens in the output.
  123. defaults to true.
  124. spaces_between_special_tokens: Whether to add spaces between special
  125. tokens in the output. Defaults to True.
  126. logits_processors: List of functions that modify logits based on
  127. previously generated tokens, and optionally prompt tokens as
  128. a first argument.
  129. truncate_prompt_tokens: If set to an integer k, will use only the last
  130. k tokens from the prompt (i.e. left-truncation). Defaults to None
  131. (i.e. no truncation).
  132. """
  133. def __init__(
  134. self,
  135. n: int = 1,
  136. best_of: Optional[int] = None,
  137. presence_penalty: float = 0.0,
  138. frequency_penalty: float = 0.0,
  139. repetition_penalty: float = 1.0,
  140. temperature: float = 1.0,
  141. temperature_last: bool = False,
  142. top_p: float = 1.0,
  143. top_k: int = -1,
  144. top_a: float = 0.0,
  145. min_p: float = 0.0,
  146. tfs: float = 1.0,
  147. eta_cutoff: float = 0.0,
  148. epsilon_cutoff: float = 0.0,
  149. typical_p: float = 1.0,
  150. smoothing_factor: float = 0.0,
  151. smoothing_curve: float = 1.0,
  152. seed: Optional[int] = None,
  153. use_beam_search: bool = False,
  154. length_penalty: float = 1.0,
  155. early_stopping: Union[bool, str] = False,
  156. stop: Union[None, str, List[str]] = None,
  157. stop_token_ids: Optional[List[int]] = None,
  158. include_stop_str_in_output: bool = False,
  159. ignore_eos: bool = False,
  160. max_tokens: Optional[int] = 16,
  161. min_tokens: int = 0,
  162. logprobs: Optional[int] = None,
  163. prompt_logprobs: Optional[int] = None,
  164. detokenize: bool = True,
  165. custom_token_bans: Optional[List[int]] = None,
  166. skip_special_tokens: bool = True,
  167. spaces_between_special_tokens: bool = True,
  168. logits_processors: Optional[List[LogitsProcessorFunc]] = None,
  169. truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None,
  170. ) -> None:
  171. self.n = n
  172. self.best_of = best_of if best_of is not None else n
  173. self.presence_penalty = presence_penalty
  174. self.frequency_penalty = frequency_penalty
  175. self.repetition_penalty = repetition_penalty
  176. self.temperature = temperature
  177. self.temperature_last = temperature_last
  178. self.top_p = top_p
  179. self.top_k = top_k
  180. self.top_a = top_a
  181. self.min_p = min_p
  182. self.tfs = tfs
  183. self.eta_cutoff = eta_cutoff
  184. self.epsilon_cutoff = epsilon_cutoff
  185. self.typical_p = typical_p
  186. self.smoothing_factor = smoothing_factor
  187. self.smoothing_curve = smoothing_curve
  188. if seed == -1:
  189. self.seed = None
  190. else:
  191. self.seed = seed
  192. self.use_beam_search = use_beam_search
  193. self.length_penalty = length_penalty
  194. self.early_stopping = early_stopping
  195. if stop is None:
  196. self.stop = []
  197. elif isinstance(stop, str):
  198. self.stop = [stop]
  199. else:
  200. self.stop = list(stop)
  201. self.stop_token_ids = stop_token_ids or []
  202. self.ignore_eos = ignore_eos
  203. self.max_tokens = max_tokens
  204. self.min_tokens = min_tokens
  205. self.logprobs = 1 if logprobs is True else logprobs
  206. self.prompt_logprobs = 1 if prompt_logprobs is True else prompt_logprobs
  207. # NOTE: This parameter is only exposed at the engine level for now.
  208. # It is not exposed in the OpenAI API server, as the OpenAI API does
  209. # not support returning only a list of token IDs.
  210. self.detokenize = detokenize
  211. self.custom_token_bans = custom_token_bans or []
  212. self.skip_special_tokens = skip_special_tokens
  213. self.spaces_between_special_tokens = spaces_between_special_tokens
  214. self.logits_processors = logits_processors or []
  215. self.include_stop_str_in_output = include_stop_str_in_output
  216. self.truncate_prompt_tokens = truncate_prompt_tokens
  217. # Number of characters to hold back for stop string evaluation
  218. # until sequence is finished.
  219. if self.stop and not include_stop_str_in_output:
  220. self.output_text_buffer_length = max(len(s) for s in self.stop) - 1
  221. else:
  222. self.output_text_buffer_length = 0
  223. self.default_values = {
  224. "n": 1,
  225. "best_of": 1,
  226. "presence_penalty": 0.0,
  227. "frequency_penalty": 0.0,
  228. "repetition_penalty": 1.0,
  229. "temperature": 1.0,
  230. "temperature_last": False,
  231. "top_p": 1.0,
  232. "top_k": -1,
  233. "top_a": 0.0,
  234. "min_p": 0.0,
  235. "tfs": 1.0,
  236. "eta_cutoff": 0.0,
  237. "epsilon_cutoff": 0.0,
  238. "typical_p": 1.0,
  239. "smoothing_factor": 0.0,
  240. "smoothing_curve": 1.0,
  241. "seed": None,
  242. "use_beam_search": False,
  243. "length_penalty": 1.0,
  244. "early_stopping": False,
  245. "stop": [],
  246. "stop_token_ids": [],
  247. "ignore_eos": False,
  248. "max_tokens": 16,
  249. "min_tokens": 0,
  250. "logprobs": None,
  251. "prompt_logprobs": None,
  252. "detokenize": True,
  253. "custom_token_bans": [],
  254. "skip_special_tokens": True,
  255. "spaces_between_special_tokens": True,
  256. "include_stop_str_in_output": False,
  257. "truncate_prompt_tokens": None,
  258. }
  259. # Number of characters to hold back for stop string evaluation
  260. # until sequence is finished.
  261. if self.stop and not include_stop_str_in_output:
  262. self.output_text_buffer_length = max(len(s) for s in self.stop) - 1
  263. else:
  264. self.output_text_buffer_length = 0
  265. self._verify_args()
  266. if self.use_beam_search:
  267. if not APHRODITE_NO_DEPRECATION_WARNING:
  268. logger.warning(
  269. "[IMPORTANT] We plan to discontinue the support for beam "
  270. "search in the next major release. Set "
  271. "APHRODITE_NO_DEPRECATION_WARNING=1 to "
  272. "suppress this warning.")
  273. self._verify_beam_search()
  274. else:
  275. self._verify_non_beam_search()
  276. if self.temperature < _SAMPLING_EPS:
  277. # Zero temperature means greedy sampling.
  278. self.top_p = 1.0
  279. self.top_k = -1
  280. self.min_p = 0.0
  281. self.top_a = 0.0
  282. self._verify_greedy_sampling()
  283. # eos_token_id is added to this by the engine
  284. self.all_stop_token_ids = set(self.stop_token_ids)
  285. def _verify_args(self) -> None:
  286. if self.n < 1:
  287. raise ValueError(f"n must be at least 1, got {self.n}.")
  288. if self.best_of < self.n:
  289. raise ValueError(f"best_of must be greater than or equal to n, "
  290. f"got n={self.n} and best_of={self.best_of}.")
  291. if not -2.0 <= self.presence_penalty <= 2.0:
  292. raise ValueError("presence_penalty must be in [-2, 2], got "
  293. f"{self.presence_penalty}.")
  294. if not -2.0 <= self.frequency_penalty <= 2.0:
  295. raise ValueError("frequency_penalty must be in [-2, 2], got "
  296. f"{self.frequency_penalty}.")
  297. if self.repetition_penalty < 1.0:
  298. raise ValueError("repetition_penalty must be in [1, inf), got "
  299. f"{self.repetition_penalty}.")
  300. if self.temperature < 0.0:
  301. raise ValueError(
  302. f"temperature must be non-negative, got {self.temperature}.")
  303. if not 0.0 < self.top_p <= 1.0:
  304. raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
  305. if self.top_k < -1 or self.top_k == 0:
  306. raise ValueError(f"top_k must be -1 (disable), or at least 1, "
  307. f"got {self.top_k}.")
  308. if self.top_a < 0:
  309. raise ValueError(f"top_a must be non negative, got {self.top_a}.")
  310. if not 0.0 <= self.min_p <= 1.0:
  311. raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.")
  312. if not 0.0 < self.tfs <= 1.0:
  313. raise ValueError(f"tfs must be in (0, 1], got {self.tfs}.")
  314. if self.epsilon_cutoff < 0.0 or self.epsilon_cutoff > 1000.0:
  315. raise ValueError("epsilon_cutoff must be in [0, 1000], got "
  316. f"{self.epsilon_cutoff}.")
  317. # pylint: disable=unneeded-not
  318. if not self.eta_cutoff >= 0:
  319. raise ValueError(
  320. f"eta_cutoff must be non negative, got {self.eta_cutoff}.")
  321. if not 0.0 <= self.typical_p <= 1.0:
  322. raise ValueError(
  323. f"typical_p must be in (0, 1], got {self.typical_p}.")
  324. if self.max_tokens is not None and self.max_tokens < 1:
  325. raise ValueError(
  326. f"max_tokens must be at least 1, got {self.max_tokens}.")
  327. if self.min_tokens < 0:
  328. raise ValueError(f"min_tokens must be greater than or equal to 0, "
  329. f"got {self.min_tokens}.")
  330. if self.max_tokens is not None and self.min_tokens > self.max_tokens:
  331. raise ValueError(
  332. f"min_tokens must be less than or equal to "
  333. f"max_tokens={self.max_tokens}, got {self.min_tokens}.")
  334. if self.logprobs is not None and self.logprobs < 0:
  335. raise ValueError(
  336. f"logprobs must be non-negative, got {self.logprobs}.")
  337. if self.prompt_logprobs is not None and self.prompt_logprobs < 0:
  338. raise ValueError("prompt_logprobs must be non-negative, got "
  339. f"{self.prompt_logprobs}.")
  340. if (self.truncate_prompt_tokens is not None
  341. and self.truncate_prompt_tokens < 1):
  342. raise ValueError(f"truncate_prompt_tokens must be >= 1, "
  343. f"got {self.truncate_prompt_tokens}")
  344. if any(not stop_str for stop_str in self.stop):
  345. raise ValueError("stop cannot contain an empty string.")
  346. if self.stop and not self.detokenize:
  347. raise ValueError(
  348. "stop strings are only supported when detokenize is True. "
  349. "Set detokenize=True to use stop.")
  350. def _verify_beam_search(self) -> None:
  351. if self.best_of == 1:
  352. raise ValueError("best_of must be greater than 1 when using beam "
  353. f"search. Got {self.best_of}.")
  354. if self.temperature > _SAMPLING_EPS:
  355. raise ValueError("temperature must be 0 when using beam search.")
  356. if self.top_p < 1.0 - _SAMPLING_EPS:
  357. raise ValueError("top_p must be 1 when using beam search.")
  358. if self.top_k != -1:
  359. raise ValueError("top_k must be -1 when using beam search.")
  360. if self.early_stopping not in [True, False, "never"]:
  361. raise ValueError(
  362. f"early_stopping must be True, False, or 'never', "
  363. f"got {self.early_stopping}.")
  364. def _verify_non_beam_search(self) -> None:
  365. if self.early_stopping is not False:
  366. raise ValueError("early_stopping is not effective and must be "
  367. "False when not using beam search.")
  368. if (self.length_penalty < 1.0 - _SAMPLING_EPS
  369. or self.length_penalty > 1.0 + _SAMPLING_EPS):
  370. raise ValueError(
  371. "length_penalty is not effective and must be the "
  372. "default value of 1.0 when not using beam search.")
  373. def _verify_greedy_sampling(self) -> None:
  374. if self.best_of > 1:
  375. raise ValueError("best_of must be 1 when using greedy sampling."
  376. f"Got {self.best_of}.")
  377. if self.top_p < 1.0 - _SAMPLING_EPS:
  378. raise ValueError("top_p must be 1 when using greedy sampling.")
  379. if self.top_k != -1:
  380. raise ValueError("top_k must be -1 when using greedy sampling.")
  381. def update_from_generation_config(
  382. self,
  383. generation_config: Dict[str, Any],
  384. model_eos_token_id: Optional[int] = None) -> None:
  385. """Update if there are non-default values from generation_config"""
  386. if model_eos_token_id is not None:
  387. # Add the eos token id into the sampling_params to support
  388. # min_tokens processing.
  389. self.all_stop_token_ids.add(model_eos_token_id)
  390. # Update eos_token_id for generation
  391. if (eos_ids := generation_config.get("eos_token_id")) is not None:
  392. # it can be either int or list of int
  393. eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids)
  394. if model_eos_token_id is not None:
  395. # We don't need to include the primary eos_token_id in
  396. # stop_token_ids since it's handled separately for stopping
  397. # purposes.
  398. eos_ids.discard(model_eos_token_id)
  399. if eos_ids:
  400. self.all_stop_token_ids.update(eos_ids)
  401. if not self.ignore_eos:
  402. eos_ids.update(self.stop_token_ids)
  403. self.stop_token_ids = list(eos_ids)
  404. @cached_property
  405. def sampling_type(self) -> SamplingType:
  406. if self.use_beam_search:
  407. return SamplingType.BEAM
  408. if self.temperature < _SAMPLING_EPS:
  409. return SamplingType.GREEDY
  410. if self.seed is not None:
  411. return SamplingType.RANDOM_SEED
  412. return SamplingType.RANDOM
  413. def clone(self) -> "SamplingParams":
  414. """Deep copy excluding LogitsProcessor objects.
  415. LogitsProcessor objects are excluded because they may contain an
  416. arbitrary, nontrivial amount of data.
  417. """
  418. logit_processor_refs = None if self.logits_processors is None else {
  419. id(lp): lp
  420. for lp in self.logits_processors
  421. }
  422. return copy.deepcopy(self, memo=logit_processor_refs)
  423. def __repr__(self) -> str:
  424. repr_str = "SamplingParams("
  425. for param, default_value in self.default_values.items():
  426. current_value = getattr(self, param)
  427. if current_value != default_value:
  428. repr_str += f"{param}={current_value}, "
  429. repr_str = repr_str.rstrip(', ') + ")"
  430. return repr_str