# 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})"
        )