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

from functools import partial
from pathlib import Path
from typing import Any

from transformers import PreTrainedTokenizerFast

class StreamingDetokenizer:
    """The streaming detokenizer interface so that we can detokenize one token at a time.

    Example usage is as follows:

        detokenizer = ...

        # Reset the tokenizer state
        detokenizer.reset()

        for token in generate(...):
            detokenizer.add_token(token.item())

            # Contains the whole text so far. Some tokens may not be included
            # since it contains whole words usually.
            detokenizer.text

            # Contains the printable segment (usually a word) since the last
            # time it was accessed
            detokenizer.last_segment

            # Contains all the tokens added so far
            detokenizer.tokens

        # Make sure that we detokenize any remaining tokens
        detokenizer.finalize()

        # Now detokenizer.text should match tokenizer.decode(detokenizer.tokens)
    """

    __slots__ = ...
    def reset(self) -> None: ...
    def add_token(self, token: int) -> None: ...
    def finalize(self) -> None: ...
    @property
    def last_segment(self) -> str:
        """Return the last segment of readable text since last time this property was accessed."""

class NaiveStreamingDetokenizer(StreamingDetokenizer):
    """NaiveStreamingDetokenizer relies on the underlying tokenizer
    implementation and should work with every tokenizer.

    Its complexity is O(T^2) where T is the longest line since it will
    repeatedly detokenize the same tokens until a new line is generated.
    """
    def __init__(self, tokenizer) -> None: ...
    def reset(self):  # -> None:
        ...
    def add_token(self, token):  # -> None:
        ...
    def finalize(self):  # -> None:
        ...
    @property
    def text(self):  # -> str:
        ...

class SPMStreamingDetokenizer(StreamingDetokenizer):
    """A streaming detokenizer for SPM models.

    It adds tokens to the text if the next token starts with the special SPM
    underscore which results in linear complexity.
    """
    def __init__(self, tokenizer, trim_space=...) -> None: ...
    def reset(self):  # -> None:
        ...
    def add_token(self, token):  # -> None:
        ...
    def finalize(self):  # -> None:
        ...

class BPEStreamingDetokenizer(StreamingDetokenizer):
    """A streaming detokenizer for OpenAI style BPE models.

    It adds tokens to the text if the next token starts with a space similar to
    the SPM detokenizer.
    """

    _byte_decoder = ...
    _space_matches = ...
    def __init__(self, tokenizer) -> None: ...
    def reset(self):  # -> None:
        ...
    def add_token(self, token):  # -> None:
        ...
    def finalize(self):  # -> None:
        ...
    @classmethod
    def make_byte_decoder(cls):  # -> None:
        """See https://github.com/openai/gpt-2/blob/master/src/encoder.py for the rationale."""

class TokenizerWrapper:
    """A wrapper that combines an HF tokenizer and a detokenizer.

    Accessing any attribute other than the ``detokenizer`` is forwarded to the
    huggingface tokenizer.
    """

    _tokenizer: PreTrainedTokenizerFast
    eos_token_id: int | None
    eos_token: str | None
    eos_token_ids: list[int] | set[int] | None
    bos_token_id: int | None
    bos_token: str | None
    vocab_size: int
    all_special_tokens: list[str]
    think_start: str | None
    think_end: str | None
    think_start_id: int | None
    think_end_id: int | None
    think_start_tokens: list[int] | None
    think_end_tokens: list[int] | None

    def __init__(
        self,
        tokenizer: Any,
        detokenizer_class: Any = ...,
        eos_token_ids: list[int] | set[int] | None = ...,
        chat_template: Any = ...,
        tool_parser: Any = ...,
        tool_call_start: str | None = ...,
        tool_call_end: str | None = ...,
    ) -> None: ...
    def encode(self, text: str, **kwargs: Any) -> list[int]: ...
    def decode(self, token_ids: list[int], **kwargs: Any) -> str: ...
    def apply_chat_template(
        self,
        messages: list[dict[str, Any]],
        tokenize: bool = False,
        add_generation_prompt: bool = False,
        tools: Any = None,
        **kwargs: Any,
    ) -> str: ...
    def get_vocab(self) -> dict[str, int]: ...
    def add_eos_token(self, token: str) -> None: ...
    @property
    def has_thinking(self) -> bool: ...
    @property
    def think_start(self) -> str | None: ...
    @property
    def think_end(self) -> str | None: ...
    @property
    def has_tool_calling(self) -> bool: ...
    @property
    def tool_call_start(self) -> str | None: ...
    @property
    def tool_call_end(self) -> str | None: ...
    @property
    def detokenizer(self) -> NaiveStreamingDetokenizer:
        """Get a stateful streaming detokenizer."""

    def __getattr__(self, attr: str) -> Any: ...
    def __setattr__(self, attr: str, value: Any) -> None: ...

class NewlineTokenizer(PreTrainedTokenizerFast):
    """A tokenizer that replaces newlines with <n> and <n> with new line."""
    def __init__(self, *args, **kwargs) -> None: ...
    def encode(self, text, **kwargs):  # -> list[int]:
        ...
    def encode_batch(self, texts, **kwargs): ...
    def decode(self, *args, **kwargs):  # -> str:
        ...
    def batch_decode(self, *args, **kwargs):  # -> list[str]:
        ...

def load(
    model_path: Path,
    tokenizer_config_extra: dict[str, Any] | None = None,
    eos_token_ids: list[int] | int | None = None,
) -> TokenizerWrapper:
    """Load a huggingface tokenizer and try to infer the type of streaming
    detokenizer to use.

    Note, to use a fast streaming tokenizer, pass a local file path rather than
    a Hugging Face repo ID.
    """

# Alias for backward compatibility
load_tokenizer = load

def no_bos_or_eos(sequence: list[int], bos: int, eos: int) -> list[int]: ...
