mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-26 16:35:05 +00:00
120 lines
5.0 KiB
Python
120 lines
5.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# Derived from vLLM; modified by DeepSeek.
|
||
|
||
from __future__ import annotations
|
||
|
||
import random
|
||
|
||
import torch
|
||
|
||
SAMPLING_EPS = 1e-5
|
||
|
||
|
||
class SamplingParams:
|
||
"""Sampling parameters for text generation.
|
||
|
||
Overall, we follow the sampling parameters from the OpenAI text completion
|
||
API (https://platform.openai.com/docs/api-reference/completions/create).
|
||
|
||
Args:
|
||
temperature: Float that controls the randomness of the sampling. Lower
|
||
values make the model more deterministic, while higher values make
|
||
the model more random. Zero means greedy sampling.
|
||
top_p: Float that controls the cumulative probability of the top tokens
|
||
to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
|
||
top_k: Integer that controls the number of top tokens to consider. Set
|
||
to -1 to consider all tokens.
|
||
ignore_eos: Whether to ignore the EOS token and continue generating
|
||
tokens after the EOS token is generated.
|
||
max_tokens: Maximum number of tokens to generate per output sequence.
|
||
logprobs: Compatibility parameter; ignored.
|
||
top_logprobs: An integer between 0 and 1000 specifying the number of most
|
||
likely tokens to return at each token position, each with an associated
|
||
log probability.
|
||
first_n_token_no_eos: Suppress EOS for the first generated token when set
|
||
to 1.
|
||
eos_id: HACK token id which will not appear in first n generated tokens.
|
||
default value is 0, which is only for warm up phase, otherwise it should
|
||
be real eos_token_id
|
||
so_args: dict, "schema_type", "schema"；可选 "trigger_token_ids"（单条 list[int]）：
|
||
触发 grammar 生效的 token id 序列，完整出现即触发，出现前自由生成；
|
||
可选 "bound_token_id"（非负整数，0 有效；缺省或 None 表示无 bound，
|
||
必须搭配非空 trigger_token_ids，列表非法）：bound 生效后不再允许自由
|
||
token，先逐 token 强制走完 trigger 序列再进 grammar
|
||
seq_seed: int(uint64), the random seed of this sequence for sampling.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
temperature: float = 1.0,
|
||
top_p: float = 1.0,
|
||
top_k: int = -1,
|
||
ignore_eos: bool = False,
|
||
max_tokens: int = 16,
|
||
logprobs: bool = False,
|
||
prompt_logprobs: bool = False,
|
||
top_logprobs: int = 0,
|
||
first_n_token_no_eos: int = 0,
|
||
eos_id: int = 0,
|
||
so_args: dict | None = None,
|
||
golden_token_ids: list[int] | None = None,
|
||
seq_seed: int | None = None,
|
||
) -> None:
|
||
del logprobs
|
||
|
||
self.temperature = temperature
|
||
self.top_p = top_p
|
||
self.top_k = top_k
|
||
self.ignore_eos = ignore_eos
|
||
self.max_tokens = max_tokens
|
||
if self.max_tokens == 0:
|
||
self.max_tokens = 1
|
||
self.prompt_logprobs = prompt_logprobs
|
||
self.top_logprobs = top_logprobs
|
||
self.seq_seed = self._make_seed(seq_seed)
|
||
self._verify_args()
|
||
if self.temperature < SAMPLING_EPS:
|
||
self.temperature = SAMPLING_EPS
|
||
self.top_p = 1.0
|
||
self.top_k = -1
|
||
self.first_n_token_no_eos = first_n_token_no_eos
|
||
self.eos_id = eos_id
|
||
self.so_args = so_args
|
||
self.golden_token_ids = golden_token_ids
|
||
|
||
@staticmethod
|
||
def _make_seed(seed: int | None) -> int:
|
||
if seed is None:
|
||
return random.getrandbits(64)
|
||
if isinstance(seed, int):
|
||
return seed & 0xFFFF_FFFF_FFFF_FFFF
|
||
raise TypeError()
|
||
|
||
def _verify_args(self) -> None:
|
||
if self.temperature < 0.0:
|
||
raise ValueError(f"temperature must be non-negative, got {self.temperature}.")
|
||
if not 0.0 < self.top_p <= 1.0:
|
||
raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
|
||
if self.top_k < -1 or self.top_k == 0 or self.top_k > torch.iinfo(torch.int32).max:
|
||
raise ValueError(f"top_k must be -1 (disable), or between 1 and int32 max, got {self.top_k}.")
|
||
if self.max_tokens < 1 or self.max_tokens > torch.iinfo(torch.int32).max:
|
||
raise ValueError(f"max_tokens must be >= 1 and in range of int32, got {self.max_tokens}.")
|
||
if self.top_logprobs < 0 or self.top_logprobs > 1000:
|
||
raise ValueError(f"top_logprobs must be in [0, 1000], got {self.top_logprobs}.")
|
||
if self.seq_seed < 0 or self.seq_seed >= 2**64:
|
||
raise ValueError(f"seq_seed must be in range of uint64_t, got {self.seq_seed}.")
|
||
|
||
def __repr__(self) -> str:
|
||
return (
|
||
f"SamplingParams("
|
||
f"temperature={self.temperature}, "
|
||
f"top_p={self.top_p}, "
|
||
f"top_k={self.top_k}, "
|
||
f"ignore_eos={self.ignore_eos}, "
|
||
f"max_tokens={self.max_tokens}, "
|
||
f"prompt_logprobs={self.prompt_logprobs}, "
|
||
f"top_logprobs={self.top_logprobs}, "
|
||
f"so_args={self.so_args}, "
|
||
f"seq_seed={self.seq_seed})"
|
||
)
|