"""
A module containing the implementation of the greedy based token proposal method,
where a single chain of tokens of length `speculative_tokens` is generated by the
speculator and verified by the verifier. This follows both the traditional
implementation of greedy as defined in the
[Spec-Bench paper](https://arxiv.org/abs/2401.07851) along with modifications to
support Spec-Verification as described in the
[SpecDec paper](https://arxiv.org/abs/2211.17192).

Classes:
    - GreedyTokenProposalConfig: Configuration for the greedy based token proposal
      method
"""

from typing import Literal

from pydantic import Field

from speculators.config import TokenProposalConfig

__all__ = ["GreedyTokenProposalConfig"]


@TokenProposalConfig.register("greedy")
class GreedyTokenProposalConfig(TokenProposalConfig):
    """
    Configuration for the greedy based token proposal method, where a single chain
    of tokens of length `speculative_tokens` is generated by the speculator and
    verified by the verifier. This follows both the traditional implementation of
    greedy as defined in the [Spec-Bench paper](https://arxiv.org/abs/2401.07851) along
    with modifications to support Spec-Verification as described in the
    [SpecDec paper](https://aclanthology.org/2023.findings-emnlp.257/).

    The default values are set to the traditional greedy implementation, where the
    top-1 token from the speculator must match the top-1 token from the verifier.
    """

    proposal_type: Literal["greedy"] = "greedy"
    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,
    )
    verifier_accept_k: int = Field(
        default=1,
        description=(
            "The top number of tokens from the verifier that the speculator's "
            "speculative token must fall within to be accepted. For 1, the "
            "speculative token must be the same as the verifier's token. "
            "For 5, the speculative token must be within the top 5 tokens."
        ),
        ge=1,
    )
    accept_tolerance: float = Field(
        default=0.0,
        description=(
            "When verifier_accept_k > 1, this is the maximum distance, "
            "in log liklihood, a speculative token found within the verifier's "
            "top k tokens can be from the verifier's top token to be accepted. "
            "This enables a limit on how likely a speculative token must be to "
            "be accepted ensuring that unlikely or poorly predicted tokens "
            "are not accepted. For 0, this is equivalent to verifier_accept_k = 1. "
            "Values up to 5 have shown to minimally impact accuracy."
        ),
        ge=0.0,
    )
