"""
This type stub file was generated by pyright.
"""

import contextlib
import mlx.core as mx
import mlx.nn as nn
from dataclasses import dataclass
from collections import deque
from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
from transformers import PreTrainedTokenizer
from .tokenizer_utils import TokenizerWrapper

DEFAULT_PROMPT = ...
DEFAULT_MAX_TOKENS = ...
DEFAULT_TEMP = ...
DEFAULT_TOP_P = ...
DEFAULT_MIN_P = ...
DEFAULT_TOP_K = ...
DEFAULT_XTC_PROBABILITY = ...
DEFAULT_XTC_THRESHOLD = ...
DEFAULT_MIN_TOKENS_TO_KEEP = ...
DEFAULT_SEED = ...
DEFAULT_MODEL = ...
DEFAULT_QUANTIZED_KV_START = ...

def str2bool(string):  # -> bool:
    ...
def setup_arg_parser():  # -> ArgumentParser:
    """Set up and return the argument parser."""
    ...

generation_stream: mx.Stream

@contextlib.contextmanager
def wired_limit(
    model: nn.Module, streams: Optional[List[mx.Stream]] = ...
):  # -> Generator[None, Any, None]:
    """
    A context manager to temporarily change the wired limit.

    Note, the wired limit should not be changed during an async eval.  If an
    async eval could be running pass in the streams to synchronize with prior
    to exiting the context manager.
    """
    ...
@dataclass
class GenerationResponse:
    """
    The output of :func:`stream_generate`.

    Args:
        text (str): The next segment of decoded text. This can be an empty string.
        token (int): The next token.
        from_draft (bool): Whether the token was generated by the draft model.
        logprobs (mx.array): A vector of log probabilities.
        prompt_tokens (int): The number of tokens in the prompt.
        prompt_tps (float): The prompt processing tokens-per-second.
        generation_tokens (int): The number of generated tokens.
        generation_tps (float): The tokens-per-second for generation.
        peak_memory (float): The peak memory used so far in GB.
        finish_reason (str): The reason the response is being sent: "length", "stop" or `None`
    """

    text: str
    token: int
    logprobs: mx.array
    from_draft: bool
    prompt_tokens: int
    prompt_tps: float
    generation_tokens: int
    generation_tps: float
    peak_memory: float
    finish_reason: Optional[str] = ...

def maybe_quantize_kv_cache(
    prompt_cache: Any,
    quantized_kv_start: int | None,
    kv_group_size: int | None,
    kv_bits: int | None,
) -> None: ...
def generate_step(
    prompt: mx.array,
    model: nn.Module,
    *,
    max_tokens: int = ...,
    sampler: Optional[Callable[[mx.array], mx.array]] = ...,
    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
    max_kv_size: Optional[int] = ...,
    prompt_cache: Optional[Any] = ...,
    prefill_step_size: int = ...,
    kv_bits: Optional[int] = ...,
    kv_group_size: int = ...,
    quantized_kv_start: int = ...,
    prompt_progress_callback: Optional[Callable[[int, int], None]] = ...,
    input_embeddings: Optional[mx.array] = ...,
) -> Generator[Tuple[mx.array, mx.array], None, None]:
    """
    A generator producing token ids based on the given prompt from the model.

    Args:
        prompt (mx.array): The input prompt.
        model (nn.Module): The model to use for generation.
        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
          generator. Default: ``256``.
        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
          token from a vector of log probabilities. Default: ``None``.
        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
          A list of functions that take tokens and logits and return the processed
          logits. Default: ``None``.
        max_kv_size (int, optional): Maximum size of the key-value cache. Old
          entries (except the first 4 tokens) will be overwritten.
        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
          provided, the cache will be updated in place.
        prefill_step_size (int): Step size for processing the prompt.
        kv_bits (int, optional): Number of bits to use for KV cache quantization.
          None implies no cache quantization. Default: ``None``.
        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
        quantized_kv_start (int): Step to begin using a quantized KV cache.
           when ``kv_bits`` is non-None. Default: ``0``.
        prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
           prompt tokens processed so far and the total number of prompt tokens.
        input_embeddings (mx.array, optional): Input embeddings to use instead of or in
          conjunction with prompt tokens. Default: ``None``.

    Yields:
        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
    """
    ...

def speculative_generate_step(
    prompt: mx.array,
    model: nn.Module,
    draft_model: nn.Module,
    *,
    num_draft_tokens: int = ...,
    max_tokens: int = ...,
    sampler: Optional[Callable[[mx.array], mx.array]] = ...,
    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
    prompt_cache: Optional[Any] = ...,
    prefill_step_size: int = ...,
    kv_bits: Optional[int] = ...,
    kv_group_size: int = ...,
    quantized_kv_start: int = ...,
) -> Generator[Tuple[mx.array, mx.array, bool], None, None]:
    """
    A generator producing token ids based on the given prompt from the model.

    Args:
        prompt (mx.array): The input prompt.
        model (nn.Module): The model to use for generation.
        draft_model (nn.Module): The draft model for speculative decoding.
        num_draft_tokens (int, optional): The number of draft tokens for
          speculative decoding. Default: ``2``.
        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
          generator. Default: ``256``.
        sampler (Callable[[mx.array], mx.array], optional): A sampler for sampling a
          token from a vector of log probabilities. Default: ``None``.
        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
          A list of functions that take tokens and logits and return the processed
          logits. Default: ``None``.
        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
          provided, the cache will be updated in place. The cache must be trimmable.
        prefill_step_size (int): Step size for processing the prompt.
        kv_bits (int, optional): Number of bits to use for KV cache quantization.
          None implies no cache quantization. Default: ``None``.
        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
        quantized_kv_start (int): Step to begin using a quantized KV cache.
           when ``kv_bits`` is non-None. Default: ``0``.

    Yields:
        Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
          and a bool indicating if the token was generated by the draft model
    """
    ...

def stream_generate(
    model: nn.Module,
    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
    prompt: Union[str, mx.array, List[int]],
    max_tokens: int = ...,
    draft_model: Optional[nn.Module] = ...,
    **kwargs: Any,
) -> Generator[GenerationResponse, None, None]:
    """
    A generator producing text based on the given prompt from the model.

    Args:
        model (nn.Module): The model to use for generation.
        tokenizer (PreTrainedTokenizer): The tokenizer.
        prompt (Union[str, mx.array, List[int]]): The input prompt string or
          integer tokens.
        max_tokens (int): The maximum number of tokens to generate.
          Default: ``256``.
        draft_model (Optional[nn.Module]): An optional draft model. If provided
          then speculative decoding is used. The draft model must use the same
          tokenizer as the main model. Default: ``None``.
        kwargs: The remaining options get passed to :func:`generate_step`.
          See :func:`generate_step` for more details.

    Yields:
        GenerationResponse: An instance containing the generated text segment and
            associated metadata. See :class:`GenerationResponse` for details.
    """
    ...

def generate(
    model: nn.Module,
    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
    prompt: Union[str, List[int]],
    verbose: bool = ...,
    **kwargs,
) -> str:
    """
    Generate a complete response from the model.

    Args:
       model (nn.Module): The language model.
       tokenizer (PreTrainedTokenizer): The tokenizer.
       prompt (Union[str, List[int]]): The input prompt string or integer tokens.
       verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
       kwargs: The remaining options get passed to :func:`stream_generate`.
          See :func:`stream_generate` for more details.
    """
    ...

def _merge_caches(caches: List[List[Any]]) -> List[Any]: ...
@dataclass
class BatchStats:
    """
    An data object to hold generation stats.

    Args:
        prompt_tokens (int): The number of prompt tokens processed.
        prompt_tps (float): The prompt processing tokens-per-second.
        prompt_time (float): The time in seconds spent in prompt processing.
        generation_tokens (int): The number of generated tokens.
        generation_tps (float): The tokens-per-second for generation.
        generation_time (float): The time in seconds spent in generation .
        peak_memory (float): The peak memory used so far in GB.
    """

    prompt_tokens: int = ...
    prompt_tps: float = ...
    prompt_time: float = ...
    generation_tokens: int = ...
    generation_tps: float = ...
    generation_time: float = ...
    peak_memory: float = ...

class SequenceStateMachine:
    """A state machine that uses one Aho-Corasick trie per state to efficiently
    track state across a generated sequence.

    The transitions are provided as state -> [(sequence, new_state)].

    Example:

        sm = SequenceStateMachine(
            transitions={
                "normal": [
                    (think_start_tokens, "reasoning"),
                    (tool_start_tokens, "tool"),
                    (eos, None),
                ],
                "reasoning": [
                    (think_end_tokens, "normal"),
                    (eos, None),
                ],
                "tool": [
                    (tool_end_tokens, None),
                    (eos, None)
                ],
            },
            initial="normal"
        )
    """
    def __init__(self, transitions=..., initial=...) -> None: ...
    def __deepcopy__(self, memo):  # -> SequenceStateMachine:
        ...
    def make_state(self):  # -> tuple[str, Any, dict[Any, Any]]:
        ...
    @staticmethod
    def match(state, x):  # -> tuple[tuple[Any, Any | None, Any], Any | None, Any]:
        ...

class PromptProcessingBatch:
    """
    A batch processor for prompt tokens with support for incremental processing.

    This class handles batched prompt processing, managing KV caches and preparing
    tokens for generation. It supports extending, filtering, and splitting batches.
    """
    @dataclass
    class Response:
        uid: int
        progress: tuple
        end_of_segment: bool
        end_of_prompt: bool
        ...

    def __init__(
        self,
        model: nn.Module,
        uids: List[int],
        caches: List[List[Any]],
        tokens: Optional[List[List[int]]] = ...,
        prefill_step_size: int = ...,
        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
        fallback_sampler: Optional[Callable[[mx.array], mx.array]] = ...,
        logits_processors: Optional[
            List[List[Callable[[mx.array, mx.array], mx.array]]]
        ] = ...,
        state_machines: Optional[List[SequenceStateMachine]] = ...,
        max_tokens: Optional[List[int]] = ...,
    ) -> None: ...
    def __len__(self):  # -> int:
        ...
    def extract_cache(self, idx: int) -> List[Any]: ...
    def extend(self, batch):  # -> None:
        ...
    def split(self, indices: List[int]):  # -> Self:
        ...
    def filter(self, keep: List[int]):  # -> None:
        ...
    def prompt(self, tokens: List[List[int]]):  # -> None:
        """
        Process prompt tokens through the model.

        Args:
            tokens: List of token sequences to process.
        """
        ...

    def generate(self, tokens: List[List[int]]):  # -> GenerationBatch:
        """
        Transition from prompt processing to generation.

        Args:
            tokens: Final tokens for each sequence to start generation.

        Returns:
            A GenerationBatch ready for token generation.
        """
        ...

    @classmethod
    def empty(
        cls,
        model: nn.Module,
        fallback_sampler: Callable[[mx.array], mx.array],
        prefill_step_size: int = ...,
    ):  # -> Self:
        ...

class GenerationBatch:
    """
    A batched token generator that manages multiple sequences in parallel.

    This class handles the generation phase after prompt processing, managing
    KV caches, sampling, and stop sequence detection for multiple sequences.
    """
    @dataclass
    class Response:
        uid: int
        token: int
        logprobs: mx.array
        finish_reason: Optional[str]
        current_state: Optional[str]
        match_sequence: Optional[List[int]]
        prompt_cache: Optional[List[Any]]
        all_tokens: Optional[List[int]]
        ...

    model: nn.Module
    uids: List[int]
    prompt_cache: List[Any]
    tokens: List[List[int]]
    samplers: Optional[List[Callable[[mx.array], mx.array]]]
    fallback_sampler: Callable[[mx.array], mx.array]
    logits_processors: Optional[List[List[Callable[[mx.array, mx.array], mx.array]]]]
    state_machines: List[SequenceStateMachine]
    max_tokens: List[int]
    _current_tokens: Optional[mx.array]
    _current_logprobs: mx.array | List[mx.array]
    _next_tokens: Optional[mx.array]
    _next_logprobs: mx.array | List[mx.array]
    _token_context: List[Any]
    _num_tokens: List[int]
    _matcher_states: List[Any]

    def __init__(
        self,
        model: nn.Module,
        uids: List[int],
        inputs: mx.array,
        prompt_cache: List[Any],
        tokens: List[List[int]],
        samplers: Optional[List[Callable[[mx.array], mx.array]]],
        fallback_sampler: Callable[[mx.array], mx.array],
        logits_processors: Optional[
            List[List[Callable[[mx.array, mx.array], mx.array]]]
        ],
        state_machines: List[SequenceStateMachine],
        max_tokens: List[int],
    ) -> None: ...
    def __len__(self) -> int: ...
    def extend(self, batch: GenerationBatch) -> None: ...
    def extract_cache(self, idx: int) -> List[Any]: ...
    def filter(self, keep: List[int]) -> None: ...
    def _step(self) -> Tuple[List[int], List[mx.array]]: ...
    def next(self) -> List[Response]:
        """
        Generate the next batch of tokens.

        Returns:
            List of Response objects for each sequence in the batch.
        """
        ...

    @classmethod
    def empty(
        cls, model: nn.Module, fallback_sampler: Callable[[mx.array], mx.array]
    ):  # -> Self:
        ...

class BatchGenerator:
    """
    A batch generator implements continuous batching.

    This class provides automatic management of prompt processing and generation
    batches, handling the transition between the two.

    It also allows for segmented prompt processing which guarantees that the
    generator will stop at these boundaries when processing an input.
    """
    def __init__(
        self,
        model: nn.Module,
        max_tokens: int = ...,
        stop_tokens: Optional[Sequence[Sequence[int]]] = ...,
        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
        logits_processors: Optional[
            List[Callable[[mx.array, mx.array], mx.array]]
        ] = ...,
        completion_batch_size: int = ...,
        prefill_batch_size: int = ...,
        prefill_step_size: int = ...,
    ) -> None: ...
    def close(self) -> None: ...
    def __del__(self):  # -> None:
        ...
    @contextlib.contextmanager
    def stats(self, stats=...):  # -> Generator[Any | BatchStats, Any, None]:
        ...
    _unprocessed_sequences: deque[tuple[Any, ...]]
    _prompt_batch: PromptProcessingBatch
    _generation_batch: GenerationBatch
    _currently_processing: list[Any]
    _gen_tokens_counter: int
    _steps_counter: int
    def _next(
        self,
    ) -> tuple[
        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
    ]: ...
    def insert(
        self,
        prompts: List[List[int]],
        max_tokens: Optional[List[int]] = ...,
        caches: Optional[List[List[Any]]] = ...,
        all_tokens: Optional[List[List[int]]] = ...,
        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
        logits_processors: Optional[
            List[List[Callable[[mx.array, mx.array], mx.array]]]
        ] = ...,
        state_machines: Optional[List[SequenceStateMachine]] = ...,
    ) -> List[int]: ...
    def insert_segments(
        self,
        segments: List[List[List[int]]],
        max_tokens: Optional[List[int]] = ...,
        caches: Optional[List[List[Any]]] = ...,
        all_tokens: Optional[List[List[int]]] = ...,
        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
        logits_processors: Optional[
            List[List[Callable[[mx.array, mx.array], mx.array]]]
        ] = ...,
        state_machines: Optional[List[SequenceStateMachine]] = ...,
    ) -> List[int]: ...
    def extract_cache(self, uids: List[int]) -> dict[int, Any]: ...
    def remove(
        self, uids: List[int], return_prompt_caches: bool = ...
    ) -> dict[int, Any]: ...
    @property
    def prompt_cache_nbytes(self) -> int: ...
    def next(
        self,
    ) -> tuple[
        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
    ]: ...
    def next_generated(self) -> List[GenerationBatch.Response]: ...

@dataclass
class BatchResponse:
    """
    A data object to hold a batch generation response.

    Args:
        texts: (List[str]): The generated text for each prompt.
        stats (BatchStats): Statistics about the generation.
    """

    texts: List[str]
    stats: BatchStats
    caches: Optional[List[List[Any]]]
    ...

def batch_generate(
    model,
    tokenizer,
    prompts: List[List[int]],
    prompt_caches: Optional[List[List[Any]]] = ...,
    max_tokens: Union[int, List[int]] = ...,
    verbose: bool = ...,
    return_prompt_caches: bool = ...,
    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
    **kwargs,
) -> BatchResponse:
    """
    Generate responses for the given batch of prompts.

    Args:
       model (nn.Module): The language model.
       tokenizer (PreTrainedTokenizer): The tokenizer.
       prompts (List[List[int]]): The input prompts.
       prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
          for each input prompt. Note, unlike ``generate_step``, the caches
          won't be updated in-place.
       verbose (bool): If ``True``, print tokens and timing information.
          Default: ``False``.
       max_tokens (Union[int, List[int]): Maximum number of output tokens. This
          can be per prompt if a list is provided.
       return_prompt_caches (bool): Return the prompt caches in the batch
          responses. Default: ``False``.
       logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
          A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
       kwargs: The remaining options get passed to :obj:`BatchGenerator`.
          See :obj:`BatchGenerator` for more details.
    """
    ...

def main():  # -> None:
    ...

if __name__ == "__main__": ...
