← back to Exo

.typings/mflux/models/common/tokenizer/tokenizer.pyi

75 lines

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

from abc import ABC, abstractmethod
from typing import Protocol, runtime_checkable
from PIL import Image
from transformers import PreTrainedTokenizer
from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput

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

@runtime_checkable
class Tokenizer(Protocol):
    tokenizer: PreTrainedTokenizer
    def tokenize(
        self,
        prompt: str | list[str],
        images: list[Image.Image] | None = ...,
        max_length: int | None = ...,
        **kwargs,
    ) -> TokenizerOutput: ...

class BaseTokenizer(ABC):
    def __init__(
        self, tokenizer: PreTrainedTokenizer, max_length: int = ...
    ) -> None: ...
    @abstractmethod
    def tokenize(
        self,
        prompt: str | list[str],
        images: list[Image.Image] | None = ...,
        max_length: int | None = ...,
        **kwargs,
    ) -> TokenizerOutput: ...

class LanguageTokenizer(BaseTokenizer):
    def __init__(
        self,
        tokenizer: PreTrainedTokenizer,
        max_length: int = ...,
        padding: str = ...,
        return_attention_mask: bool = ...,
        template: str | None = ...,
        use_chat_template: bool = ...,
        chat_template_kwargs: dict | None = ...,
        add_special_tokens: bool = ...,
    ) -> None: ...
    def tokenize(
        self,
        prompt: str | list[str],
        images: list[Image.Image] | None = ...,
        max_length: int | None = ...,
        **kwargs,
    ) -> TokenizerOutput: ...

class VisionLanguageTokenizer(BaseTokenizer):
    def __init__(
        self,
        tokenizer: PreTrainedTokenizer,
        processor,
        max_length: int = ...,
        template: str | None = ...,
        image_token: str = ...,
    ) -> None: ...
    def tokenize(
        self,
        prompt: str | list[str],
        images: list[Image.Image] | None = ...,
        max_length: int | None = ...,
        **kwargs,
    ) -> TokenizerOutput: ...