"""Approximate Speculative Decoding (ASD) as a Transformers custom generation method.

ASD is a training-free, deterministic acceptance policy for *greedy* assisted
decoding. Instead of discarding the whole draft suffix at the first position
where the draft token differs from the target argmax, it accepts a draft prefix
while ALL of the following hold:

1. request-level cumulative regret ``sum(r_j) <= assistant_asd_budget``;
2. local regret ``r_i / (K - i) <= assistant_asd_local_ratio``
   (suffix-value weighting: later draft positions get less slack);
3. relaxed mismatches in the current block ``<= assistant_asd_max_mismatches``.

Here ``r_i = max_v z_i(v) - z_i(x_i)`` is the local regret of draft token
``x_i`` under target logits ``z_i`` (``r_i = 0`` iff ``x_i`` is the target
argmax) and ``K`` is the candidate length. The policy uses *target logits only*
(no draft logits, no RNG) and exactly reduces to strict greedy verification
when the budget, the local ratio, or the mismatch cap is ``0``.

Paper: "ASD: Approximate Speculative Decoding", https://arxiv.org/abs/2608.03447

Implementation notes
--------------------
`generate` below mirrors the signature of `GenerationMixin.generate` (this is
the contract for Hub custom generation methods). It validates the ASD knobs and
then temporarily swaps `type(model)._assisted_decoding` for the ASD loop in
this file, so the standard `generate()` pipeline (logits processors, stopping
criteria, cache preparation, streamer, assistant handling, ...) runs unchanged.
The swap is reverted in a `finally` block.

`_asd_assisted_decoding` is a byte-faithful copy of
`transformers.generation.utils.GenerationMixin._assisted_decoding`
(transformers v5.19.0) with exactly three changes, marked `ASD-CHANGE`:
  1. request-level risk state is initialised before the decode loop;
  2. the strict greedy prefix match (upstream "Case 2") is replaced by the ASD
     acceptance rule (the sampling path is rejected in `generate`, so upstream
     "Case 1" is unreachable here);
  3. the risk account is charged after the EOS/length crop, only for draft
     tokens that were actually committed.
"""

from __future__ import annotations

import copy
import math
from typing import TYPE_CHECKING, Any, NamedTuple

import torch
from transformers import GenerationMixin
from transformers.cache_utils import StaticCache
from transformers.generation.candidate_generator import (
    AssistedCandidateGenerator,
    MTPCandidateGenerator,
    _prepare_attention_mask,
    _prepare_position_ids,
    _prepare_token_type_ids,
)
from transformers.generation.utils import (
    ALL_CACHE_NAMES,
    GenerateDecoderOnlyOutput,
    GenerateEncoderDecoderOutput,
    _split_model_outputs,
)

if TYPE_CHECKING:
    from transformers import (
        GenerationConfig,
        LogitsProcessorList,
        PreTrainedModel,
        PreTrainedTokenizerBase,
        StoppingCriteriaList,
    )
    from transformers.generation import BaseStreamer
    from transformers.generation.utils import GenerateNonBeamOutput


class _ASDVerification(NamedTuple):
    """Resolved ASD acceptance parameters for one request."""

    budget: float  # request-level cumulative regret budget B (inf = uncapped)
    local_ratio: float  # per-token regret / suffix-value cap g (inf = uncapped)
    max_mismatches: int  # per-block relaxed mismatch cap m
    strict_identity: bool  # True => exactly upstream strict argmax verification


_ASD_CONFIG_ATTR = "_asd_verification"


def _resolve_asd_params(
    budget: Any,
    local_ratio: Any,
    max_mismatches: Any,
) -> _ASDVerification:
    """Validate the user-facing knobs and return the resolved acceptance config.

    All three knobs ``None`` means "ASD disabled": the policy is strict argmax
    verification, i.e. token-identical to native assisted decoding. When any
    knob is set, the unset ones default to their permissive limit. A value of
    ``0`` for any knob forces strict verification (the paper's correctness
    baseline).
    """
    all_none = budget is None and local_ratio is None and max_mismatches is None

    if budget is not None:
        budget = float(budget)
        if not math.isfinite(budget) or budget < 0:
            raise ValueError(f"`assistant_asd_budget` must be a finite float >= 0, got {budget}.")
    else:
        budget = math.inf

    if local_ratio is not None:
        local_ratio = float(local_ratio)
        if not math.isfinite(local_ratio) or local_ratio < 0:
            raise ValueError(
                f"`assistant_asd_local_ratio` must be a finite float >= 0, got {local_ratio}."
            )
    else:
        local_ratio = math.inf

    if max_mismatches is not None:
        max_mismatches = int(max_mismatches)
        if max_mismatches < 0:
            raise ValueError(f"`assistant_asd_max_mismatches` must be an int >= 0, got {max_mismatches}.")
    else:
        max_mismatches = 2**31 - 1

    strict_identity = all_none or budget == 0.0 or local_ratio == 0.0 or max_mismatches == 0
    return _ASDVerification(budget, local_ratio, max_mismatches, strict_identity)


def _effective_generation_value(name: str, kwargs: dict, generation_config, model):
    """Resolve a generation flag the way `generate()` would: kwargs > config arg > model config."""
    if name in kwargs:
        return kwargs[name]
    if generation_config is not None and getattr(generation_config, name, None) is not None:
        return getattr(generation_config, name)
    return getattr(model.generation_config, name, None)


def _choose_asd_prefix(
    candidate_new_tokens: torch.LongTensor,
    new_logits: torch.FloatTensor,
    target_argmax: torch.LongTensor,
    config: _ASDVerification,
) -> tuple[int, list[float], list[bool]]:
    """Return ``(accepted, regrets, mismatches)`` for one draft block.

    ``candidate_new_tokens`` is ``(1, K)``, ``new_logits`` is ``(1, K, V)`` (the
    processed target logits at each drafted position) and ``target_argmax`` is
    ``(1, K)``: the target argmax token per position, computed with the exact
    same ``argmax`` op upstream strict verification uses, so the strict path is
    token-identical by construction. ``accepted`` is the longest prefix
    satisfying all ASD constraints; ``regrets``/``mismatches`` are the
    per-position values so the caller can charge the budget after the
    EOS/length crop, only for tokens that were actually committed.
    """
    candidate_length = candidate_new_tokens.shape[1]
    if candidate_length == 0:
        return 0, [], []

    # r_i = max_v z_i(v) - z_i(x_i); mismatch_i = (x_i != argmax_v z_i(v))
    top_logits = new_logits[0].max(dim=-1).values
    draft_logits = new_logits[0].gather(-1, candidate_new_tokens[0].unsqueeze(-1)).squeeze(-1)
    regrets = (top_logits - draft_logits).clamp_min(0.0).tolist()
    mismatches = (candidate_new_tokens[0] != target_argmax[0]).tolist()

    # `config.budget` is the *remaining* request budget for this block: the caller
    # (`_asd_assisted_decoding`) subtracts the running spend before calling.
    remaining = max(0.0, config.budget)
    cumulative_risk = 0.0
    cumulative_relaxed = 0
    accepted = 0
    for index in range(candidate_length):
        mismatch = mismatches[index]
        regret = regrets[index]
        value = float(candidate_length - index)  # suffix-value weighting

        if mismatch:
            if config.strict_identity:
                break
            cumulative_risk += regret
            cumulative_relaxed += 1

        worthwhile = (not mismatch) or (regret / value <= config.local_ratio)
        affordable = cumulative_risk <= remaining + 1e-9
        within_cap = cumulative_relaxed <= config.max_mismatches

        if not (worthwhile and affordable and within_cap):
            break
        accepted = index + 1

    return accepted, regrets, mismatches


def _asd_assisted_decoding(
    self: "PreTrainedModel",
    input_ids: torch.LongTensor,
    logits_processor: "LogitsProcessorList",
    stopping_criteria: "StoppingCriteriaList",
    generation_config: "GenerationConfig",
    synced_gpus: bool = False,
    streamer: "BaseStreamer | None" = None,
    inputs_tensor: torch.FloatTensor | None = None,
    assistant_model: "PreTrainedModel | None" = None,
    assistant_tokenizer: "PreTrainedTokenizerBase | None" = None,
    tokenizer: "PreTrainedTokenizerBase | None" = None,
    **model_kwargs,
) -> "GenerateNonBeamOutput | torch.LongTensor":
    """Drop-in replacement for `GenerationMixin._assisted_decoding` (greedy only).

    Identical to upstream except for the ASD changes marked `ASD-CHANGE`.
    """
    asd: _ASDVerification | None = getattr(self, _ASD_CONFIG_ATTR, None)
    if asd is None:  # defensive: behave exactly like upstream strict verification
        asd = _ASDVerification(math.inf, math.inf, 2**31 - 1, True)

    # The cache must be dynamic for assisted generation, and the check must happen AFTER preparing cache
    if not model_kwargs["use_cache"]:
        raise ValueError("assisted generate requires `use_cache=True`")
    if (
        generation_config.cache_implementation in ["static", "hybrid", "sliding_window"]
        or type(model_kwargs.get("past_key_values")) is StaticCache
    ):
        raise ValueError("assisted generate is not supported with Static cache classes`")

    # Same tensor the stopping criteria are built from
    eos_token_id = getattr(generation_config, "_eos_token_tensor", None)

    # Make sure we can record past on the cache
    cache = model_kwargs.get("past_key_values")
    if cache is None:
        raise RuntimeError("assisted decoding requires a cache")
    cache.activate_past_recording()

    # Get the candidate generator, given the parameterization
    candidate_generator = self._get_candidate_generator(
        generation_config=generation_config,
        input_ids=input_ids,
        inputs_tensor=inputs_tensor,
        assistant_model=assistant_model,
        logits_processor=logits_processor,
        target_tokenizer=tokenizer,
        assistant_tokenizer=assistant_tokenizer,
        model_kwargs=model_kwargs,
    )
    # init values
    output_attentions = generation_config.output_attentions
    output_hidden_states = generation_config.output_hidden_states
    output_scores = generation_config.output_scores
    output_logits = generation_config.output_logits
    return_dict_in_generate = generation_config.return_dict_in_generate

    # init attention / hidden states / scores tuples
    scores = () if (return_dict_in_generate and output_scores) else None
    raw_logits = () if (return_dict_in_generate and output_logits) else None
    decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
    cross_attentions = () if (return_dict_in_generate and output_attentions) else None
    decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None

    # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
    if return_dict_in_generate and self.config.is_encoder_decoder:
        encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
        encoder_hidden_states = (
            model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
        )

    # keep track of which sequences are already finished
    batch_size, cur_len = input_ids.shape[:2]
    if batch_size > 1:
        raise ValueError("assisted generate is only supported for batch_size = 1")
    unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)

    this_peer_finished = False
    is_first_iteration = True  # to preserve the same API in the output as other generation methods
    outputs = None
    n_matches = 0
    total_risk_spent = 0.0  # ASD-CHANGE (1): request-level cumulative regret account
    while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
        cur_len = input_ids.shape[1]

        #  1. Fetch candidate sequences from a `CandidateGenerator` and move to the correct device
        if candidate_generator.requires_model_outputs:
            candidate_input_ids, candidate_logits = candidate_generator.get_candidates(
                input_ids,
                model_kwargs=model_kwargs,
                model_outputs=outputs,
                is_first_iteration=is_first_iteration,
                n_last_matches=n_matches,
            )
        else:
            candidate_input_ids, candidate_logits = candidate_generator.get_candidates(input_ids)
        candidate_input_ids = candidate_input_ids.to(self.device)
        if candidate_logits is not None:
            candidate_logits = candidate_logits.to(self.device)

        candidate_length = candidate_input_ids.shape[1] - input_ids.shape[1]

        # 2. Use the original model to obtain the next token logits given the candidate sequence. We obtain
        # `candidate_length + 1` relevant logits from this process: in the event that all candidates are correct,
        # we use this forward pass to also pick the subsequent logits in the original model.

        # 2.1. Prepare the model inputs
        candidate_kwargs = copy.copy(model_kwargs)
        candidate_kwargs = _prepare_attention_mask(
            candidate_kwargs, candidate_input_ids.shape[1], self.config.is_encoder_decoder
        )
        candidate_kwargs = _prepare_token_type_ids(candidate_kwargs, candidate_input_ids.shape[1])
        if (position_ids := candidate_kwargs.get("position_ids")) is not None and candidate_length > 0:
            new_length = candidate_length + position_ids.shape[-1]
            candidate_kwargs = _prepare_position_ids(candidate_kwargs, new_length, self.config.is_encoder_decoder)

        next_sequence_length = candidate_length + 1 if not is_first_iteration else None
        model_inputs = self.prepare_inputs_for_generation(
            candidate_input_ids,
            next_sequence_length=next_sequence_length,
            is_first_iteration=is_first_iteration,
            **candidate_kwargs,
        )

        if "logits_to_keep" in model_inputs:
            model_inputs["logits_to_keep"] = candidate_length + 1

        # We usually need to force returning hidden_states and some others if the candidate_generator requires the model's outputs
        if candidate_generator.requires_model_outputs:
            model_inputs |= candidate_generator.model_kwargs_overrides

        # 2.2. Run a forward pass on the candidate sequence
        outputs = self(**model_inputs)

        # 2.3. Process the new logits
        # .float() is needed to retain precision for later logits manipulations
        new_logits = outputs.logits[:, -candidate_length - 1 :].to(
            dtype=torch.float32, device=input_ids.device
        )  # excludes the input prompt if present
        next_token_logits = new_logits.clone()
        if len(logits_processor) > 0:
            for i in range(candidate_length + 1):
                new_logits[:, i, :] = logits_processor(candidate_input_ids[:, : cur_len + i], new_logits[:, i, :])

        # 3. Select the accepted tokens.
        # ASD-CHANGE (2): the greedy verification below replaces upstream's strict argmax
        # prefix match with the ASD acceptance rule. The sampling path ("Case 1") and the
        # `assistant_ensemble_weight` variant are unreachable: `generate()` rejects
        # `do_sample=True` and `assistant_ensemble_weight` before installing this method.
        candidate_new_tokens = candidate_input_ids[:, cur_len:]
        # Same op as upstream strict verification, so strict mode matches by construction.
        target_argmax = new_logits.argmax(dim=-1)  # (1, candidate_length + 1)
        block_config = _ASDVerification(
            budget=asd.budget - total_risk_spent,  # remaining request budget for this block
            local_ratio=asd.local_ratio,
            max_mismatches=asd.max_mismatches,
            strict_identity=asd.strict_identity,
        )
        accepted, regrets, mismatches = _choose_asd_prefix(
            candidate_new_tokens, new_logits[:, :-1, :], target_argmax[:, :-1], block_config
        )

        # Commit the accepted draft prefix, then the target argmax at the first position
        # after it as the bonus token (deterministic, exactly as upstream: upstream's
        # `selected_tokens[:, :n_matches + 1]` is the argmax at position `n_matches`).
        bonus_token = target_argmax[:, accepted : accepted + 1]
        if accepted > 0:
            valid_tokens = torch.cat([candidate_new_tokens[:, :accepted], bonus_token], dim=-1)
        else:
            valid_tokens = bonus_token
        n_matches = accepted

        # Whenever we are drafting several tokens at once with the candidate without capping its draft (e.g. MTP), we need to
        # make sure that we did not just validate tokens outside the max length, or outside an eos token
        # Note that we should technically crop based on other stopping criteria as well in all generality, not only EOS and Length
        tokens_budget = generation_config.max_length - input_ids.shape[1]
        # This is for the max length
        if valid_tokens.shape[1] > tokens_budget:
            valid_tokens = valid_tokens[:, :tokens_budget]
        # This is for eos tokens
        if eos_token_id is not None:
            # We are restricted to batch_size == 1, so we can squeeze the batch dim to simplify
            eos_positions = torch.isin(valid_tokens.squeeze(0), eos_token_id.to(valid_tokens.device)).nonzero()
            if eos_positions.numel() > 0:
                num_drafted = eos_positions[0].item() + 1
                valid_tokens = valid_tokens[:, :num_drafted]
        # Recompute how many matches we accepted if we just cropped above due to length/eos
        n_matches = valid_tokens.shape[1] - 1

        # ASD-CHANGE (3): charge the request-level regret account after the crop, only
        # for relaxed (mismatched) draft tokens that were actually committed.
        for index in range(max(0, min(n_matches, accepted))):
            if mismatches[index]:
                total_risk_spent += regrets[index]

        # 4. Update variables according to the number of matching assistant tokens. Remember: the token generated
        # by the model after the last candidate match is also valid, as it is generated from a correct sequence.
        # Because of this last token, assisted generation search reduces to a normal greedy search/sample if there
        # is no match.

        # 4.1. Get the valid continuation, after the matching tokens
        input_ids = torch.cat((input_ids, valid_tokens), dim=-1)
        if streamer is not None:
            streamer.put(valid_tokens.cpu())
        new_cur_len = input_ids.shape[1]

        # 4.2. Discard past key values relative to unused assistant tokens if any - if `number_of_tokens_to_crop == 0`, `crop` will
        # still take care of shrinking back sliding window or linear attention cache layers back to their max size, so it's important
        # to call it still
        number_of_tokens_to_crop = candidate_length - n_matches
        outputs.past_key_values.crop(-number_of_tokens_to_crop)

        # 5. Update the candidate generation strategy if needed
        candidate_generator.update_candidate_strategy(input_ids, new_logits, n_matches)

        # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
        model_kwargs = self._update_model_kwargs_for_generation(
            outputs,
            model_kwargs,
            is_encoder_decoder=self.config.is_encoder_decoder,
            num_new_tokens=n_matches + 1,
        )
        if synced_gpus and this_peer_finished:
            continue

        # Store scores, attentions and hidden_states when required
        # Assistant: modified to append one tuple element per token, as in the other generation methods.
        if return_dict_in_generate:
            newly_added_length = n_matches + 1
            if output_scores:
                scores += tuple(new_logits[:, i, :] for i in range(newly_added_length))
            if output_logits:
                raw_logits += tuple(next_token_logits[:, i, :] for i in range(newly_added_length))

            newly_added_length = new_cur_len if is_first_iteration else newly_added_length
            if output_attentions:
                if self.config.is_encoder_decoder:
                    cross_attentions = _split_model_outputs(
                        cross_attentions, outputs.cross_attentions, cur_len, newly_added_length
                    )
                    decoder_attentions = _split_model_outputs(
                        decoder_attentions,
                        outputs.decoder_attentions,
                        cur_len,
                        newly_added_length,
                        is_decoder_attention=True,
                    )
                # some (V)LLMs have hard requirement on SDPA and thus never return attn
                elif outputs.attentions[0] is not None:
                    decoder_attentions = _split_model_outputs(
                        decoder_attentions,
                        outputs.attentions,
                        cur_len,
                        newly_added_length,
                        is_decoder_attention=True,
                    )
            if output_hidden_states:
                if self.config.is_encoder_decoder:
                    decoder_hidden_states = _split_model_outputs(
                        decoder_hidden_states, outputs.decoder_hidden_states, cur_len, newly_added_length
                    )
                else:
                    decoder_hidden_states = _split_model_outputs(
                        decoder_hidden_states, outputs.hidden_states, cur_len, newly_added_length
                    )

        unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)
        this_peer_finished = unfinished_sequences.max() == 0
        is_first_iteration = False

    if streamer is not None:
        streamer.end()

    if (
        isinstance(candidate_generator, AssistedCandidateGenerator)
        and not isinstance(candidate_generator, MTPCandidateGenerator)
        and candidate_generator.assistant_model.generation_config.num_assistant_tokens_schedule == "heuristic"
    ):
        candidate_generator.assistant_model.generation_config.num_assistant_tokens = (
            candidate_generator.num_assistant_tokens
        )
    if return_dict_in_generate:
        cache = None
        if any(cache_key in model_kwargs for cache_key in ALL_CACHE_NAMES):
            cache_key = next(cache_key for cache_key in ALL_CACHE_NAMES if cache_key in model_kwargs)
            cache = model_kwargs[cache_key]
        if self.config.is_encoder_decoder:
            return GenerateEncoderDecoderOutput(
                sequences=input_ids,
                scores=scores,
                logits=raw_logits,
                encoder_attentions=encoder_attentions,
                encoder_hidden_states=encoder_hidden_states,
                decoder_attentions=decoder_attentions,
                cross_attentions=cross_attentions,
                decoder_hidden_states=decoder_hidden_states,
                past_key_values=cache,
            )
        else:
            return GenerateDecoderOnlyOutput(
                sequences=input_ids,
                scores=scores,
                logits=raw_logits,
                attentions=decoder_attentions,
                hidden_states=decoder_hidden_states,
                past_key_values=cache,
            )
    else:
        return input_ids


def generate(
    model: "PreTrainedModel",
    inputs: torch.Tensor | None = None,
    generation_config: "GenerationConfig | None" = None,
    logits_processor: "LogitsProcessorList | None" = None,
    stopping_criteria: "StoppingCriteriaList | None" = None,
    prefix_allowed_tokens_fn=None,
    synced_gpus: bool | None = None,
    assistant_model: "PreTrainedModel | None" = None,
    streamer: "BaseStreamer | None" = None,
    negative_prompt_ids: torch.Tensor | None = None,
    negative_prompt_attention_mask: torch.Tensor | None = None,
    assistant_asd_budget: float | None = None,
    assistant_asd_local_ratio: float | None = None,
    assistant_asd_max_mismatches: int | None = None,
    **kwargs,
):
    """Custom generation entry point loaded from the Hub.

    Accepts every argument of `GenerationMixin.generate` plus the three ASD
    knobs below, then runs the standard `generate()` pipeline with the ASD
    assisted-decoding loop installed. With no ASD knob set (or any knob equal
    to ``0``), the result is token-identical to native greedy assisted decoding.

    Parameters
    ----------
    assistant_asd_budget:
        Request-level cumulative regret budget ``B >= 0``. ``None`` (default)
        leaves the budget uncapped; ``0.0`` forces strict verification.
    assistant_asd_local_ratio:
        Per-token cap ``g >= 0`` on ``r_i / (K - i)``. ``None`` leaves it
        uncapped; ``0.0`` forces strict verification.
    assistant_asd_max_mismatches:
        Per-block cap ``m >= 0`` on relaxed mismatches. ``None`` leaves it
        uncapped; ``0`` forces strict verification.
    """
    asd = _resolve_asd_params(assistant_asd_budget, assistant_asd_local_ratio, assistant_asd_max_mismatches)
    any_asd_knob = any(
        knob is not None
        for knob in (assistant_asd_budget, assistant_asd_local_ratio, assistant_asd_max_mismatches)
    )

    if _effective_generation_value("do_sample", kwargs, generation_config, model) is True:
        raise ValueError(
            "ASD is a deterministic verifier for greedy decoding (`do_sample=False`). "
            "For the sampling path, use native assisted decoding with `assistant_ensemble_weight` instead."
        )
    if _effective_generation_value("assistant_ensemble_weight", kwargs, generation_config, model) is not None:
        raise ValueError(
            "`assistant_ensemble_weight` and ASD are mutually exclusive verification relaxations. "
            "Please use only one of them."
        )
    if (
        any_asd_knob
        and assistant_model is None
        and _effective_generation_value("prompt_lookup_num_tokens", kwargs, generation_config, model) is None
    ):
        raise ValueError(
            "ASD knobs require a draft source: pass `assistant_model` "
            "(or set `prompt_lookup_num_tokens` for prompt-lookup drafting)."
        )

    # Install the ASD loop as the model class's assisted-decoding method for the
    # duration of this call. `generate()` dispatches assisted generation through
    # `getattr(type(self), "_assisted_decoding")`, so a class-level swap is the
    # supported seam; it is reverted in `finally` below.
    cls = type(model)
    was_inherited = "_assisted_decoding" not in cls.__dict__
    original = cls._assisted_decoding
    cls._assisted_decoding = _asd_assisted_decoding
    setattr(model, _ASD_CONFIG_ATTR, asd)
    try:
        return GenerationMixin.generate(
            model,
            inputs=inputs,
            generation_config=generation_config,
            logits_processor=logits_processor,
            stopping_criteria=stopping_criteria,
            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
            synced_gpus=synced_gpus,
            assistant_model=assistant_model,
            streamer=streamer,
            negative_prompt_ids=negative_prompt_ids,
            negative_prompt_attention_mask=negative_prompt_attention_mask,
            **kwargs,
        )
    finally:
        if was_inherited:
            del cls._assisted_decoding
        else:
            cls._assisted_decoding = original
        try:
            delattr(model, _ASD_CONFIG_ATTR)
        except AttributeError:
            pass