"""
A module containing the implementation of the sampling based token proposal method,
where a single chain of tokens of length `speculative_tokens` is generated by the
speculator and verified by the verifier. The scoring of the tokens is done by
comparing the speculator's top token probability, q_i, with the verifier's top
token probability, p_i. If q_i > p_i, the token is accepted. If q_i < p_i, it is
rejected with probability 1 - p_i / q_i. This is described in more detail in the
[SpecBench paper](https://arxiv.org/abs/2401.07851) and the
[Fast Inference from Transformers via Speculative Decoding paper](https://arxiv.org/abs/2211.17192).

Classes:
    - SamplingTokenProposalConfig: Configuration for the sampling based token proposal
      method
"""

from typing import Literal

from pydantic import Field

from speculators.config import TokenProposalConfig

__all__ = ["SamplingTokenProposalConfig"]


@TokenProposalConfig.register("sampling")
class SamplingTokenProposalConfig(TokenProposalConfig):
    """
    Configuration for the sampling based token proposal method, where a single chain
    of tokens of length `speculative_tokens` is generated by the speculator and
    verified by the verifier. The scoring of the tokens is done by comparing the
    speculator's top token probability, q_i, with the verifier's top token probability,
    p_i. If q_i > p_i, the token is accepted. If q_i < p_i, it is rejected with
    probability 1 - p_i / q_i. This is described in more detail in the
    [SpecBench paper](https://arxiv.org/abs/2401.07851). Additionally, a further
    lenience parameter can be applied to the verifier's probability to increase the
    acceptance likelihood. This is described in more detail in the
    [Fast Inference from Transformers via Speculative Decoding paper](https://arxiv.org/abs/2211.17192).

    The default values are set to follow the traditional sampling implementation, where
    lenience is set to 1.0, meaning the verifier's probability is unchanged.
    """

    proposal_type: Literal["sampling"] = Field(
        default="sampling",
        description="The type of this token proposal.",
    )
    speculative_tokens: int = Field(
        default=5,
        description=(
            "The number of tokens created by the speculator to run through the "
            "verifier on each forward pass. This is the maximum number of tokens "
            "that can be accepted by the verifier. Larger values increase the "
            "number of tokens that can be accepted, and therefore the potential "
            "length of accepted tokens, but also increases the computation time "
            "for both the speculator and verifier."
        ),
        ge=1,
    )
    accept_lenience: float = Field(
        default=1.0,
        description=(
            "A parameter that increases the acceptance likelihood by multiplying the "
            "verifier's probability of a speculative token by this value. "
            "1.0 (default) means the verifier's probability is unchanged, 0.5 means "
            "the verifier's probability is halved, and 0.0 means the verifier's "
            "probability is ignored. Must be between 0.0 and 1.0."
        ),
        ge=0.0,
        le=1.0,
    )
