




"""simple docstring"""




import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_gpta import GPTaTokenizer


if TYPE_CHECKING:
			from transformers.pipelines.conversational import Conversation


a_        =       logging.get_logger(__name__)

a_        =       {"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

a_        =       {
    """vocab_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/vocab.json""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/vocab.json""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/vocab.json""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/vocab.json""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/merges.txt""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/merges.txt""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/merges.txt""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/merges.txt""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/tokenizer.json""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/tokenizer.json""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/tokenizer.json""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/tokenizer.json""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/tokenizer.json""",
    },
}

a_        =       {
    """gpt2""": 1_0_2_4,
    """gpt2-medium""": 1_0_2_4,
    """gpt2-large""": 1_0_2_4,
    """gpt2-xl""": 1_0_2_4,
    """distilgpt2""": 1_0_2_4,
}

class   UpperCAmelCase_    (			__SCREAMING_SNAKE_CASE		):
		UpperCamelCase							=VOCAB_FILES_NAMES
		UpperCamelCase							=PRETRAINED_VOCAB_FILES_MAP
		UpperCamelCase							=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		UpperCamelCase							=["""input_ids""", """attention_mask"""]
		UpperCamelCase							=GPTaTokenizer
		def __init__(						self				,						UpperCamelCase_=None				,						UpperCamelCase_=None				,						UpperCamelCase_=None				,						UpperCamelCase_="<|endoftext|>"				,						UpperCamelCase_="<|endoftext|>"				,						UpperCamelCase_="<|endoftext|>"				,						UpperCamelCase_=False				,						**UpperCamelCase_				,						)						->       int:
					super().__init__(
					    A_				,						A_				,						tokenizer_file=A_				,						unk_token=A_				,						bos_token=A_				,						eos_token=A_				,						add_prefix_space=A_				,						**A_				,						)

					__lowercase			:  Union[str, Any]								=       kwargs.pop('''add_bos_token'''				,						A_       )

					__lowercase			:  Dict								=       json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()       )
					if pre_tok_state.get('''add_prefix_space'''				,						A_       ) != add_prefix_space:
								__lowercase			:  Optional[Any]								=       getattr(A_				,						pre_tok_state.pop('''type'''       )       )
								__lowercase			:  Any								=       add_prefix_space
								__lowercase			:  int								=       pre_tok_class(**A_       )

					__lowercase			:  List[Any]								=       add_prefix_space
		def 	_lowerCamelCase							(						self				,						*UpperCamelCase_				,						**UpperCamelCase_       )						->       int:
					__lowercase			:  int								=       kwargs.get('''is_split_into_words'''				,						A_       )
					assert self.add_prefix_space or not is_split_into_words, (
					    F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
					    "to use it with pretokenized inputs."
					)

					return super()._batch_encode_plus(*A_				,						**A_       )
		def 	_lowerCamelCase							(						self				,						*UpperCamelCase_				,						**UpperCamelCase_       )						->       int:
					__lowercase			:  Any								=       kwargs.get('''is_split_into_words'''				,						A_       )

					assert self.add_prefix_space or not is_split_into_words, (
					    F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
					    "to use it with pretokenized inputs."
					)

					return super()._encode_plus(*A_				,						**A_       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_ = None       )						->       str:
					__lowercase			:  List[str]								=       self._tokenizer.model.save(A_				,						name=A_       )
					return tuple(A_       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_       )						->       List[str]:
					__lowercase			:  Union[str, Any]								=       []
					for is_user, text in conversation.iter_texts():
								input_ids.extend(self.encode(A_				,						add_special_tokens=A_       ) + [self.eos_token_id]       )

					if len(A_       ) > self.model_max_length:
								__lowercase			:  Any								=       input_ids[-self.model_max_length :]
					return input_ids










import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
_A					:				Optional[int]				=			object()

# For specifying empty leaf dict `{}`
_A					:				Tuple				=			object()


def 							__snake_case							(    lowerCAmelCase_ ,							lowerCAmelCase_	)				->						Dict:
    SCREAMING_SNAKE_CASE__											=			tuple((re.compile(x + '''$'''	) for x in qs)	)
    for i in range(len(lowerCAmelCase_	) - len(lowerCAmelCase_	) + 1	):
        SCREAMING_SNAKE_CASE__											=			[x.match(lowerCAmelCase_	) for x, y in zip(lowerCAmelCase_ ,							ks[i:]	)]
        if matches and all(lowerCAmelCase_	):
            return True
    return False


def 							__snake_case							(    lowerCAmelCase_	)				->						int:
    def replace(lowerCAmelCase_ ,							lowerCAmelCase_	):
        for rule, replacement in rules:
            if _match(lowerCAmelCase_ ,							lowerCAmelCase_	):
                return replacement
        return val

    return replace


def 							__snake_case							(    )				->						Optional[int]:
    return [
        # embeddings
        (("transformer", "wpe", "embedding"), P('''mp''' ,							lowerCAmelCase_	)),
        (("transformer", "wte", "embedding"), P('''mp''' ,							lowerCAmelCase_	)),
        # atention
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(lowerCAmelCase_ ,							'''mp'''	)),
        (("attention", "out_proj", "kernel"), P('''mp''' ,							lowerCAmelCase_	)),
        (("attention", "out_proj", "bias"), None),
        # mlp
        (("mlp", "c_fc", "kernel"), P(lowerCAmelCase_ ,							'''mp'''	)),
        (("mlp", "c_fc", "bias"), P('''mp'''	)),
        (("mlp", "c_proj", "kernel"), P('''mp''' ,							lowerCAmelCase_	)),
        (("mlp", "c_proj", "bias"), None),
        # layer norms
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),
    ]



def 							__snake_case							(    lowerCAmelCase_	)				->						Optional[Any]:
    SCREAMING_SNAKE_CASE__											=			_get_partition_rules()
    SCREAMING_SNAKE_CASE__											=			_replacement_rules(lowerCAmelCase_	)
    SCREAMING_SNAKE_CASE__											=			{k: _unmatched for k in flatten_dict(lowerCAmelCase_	)}
    SCREAMING_SNAKE_CASE__											=			{k: replace(lowerCAmelCase_ ,							lowerCAmelCase_	) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."
    return freeze(unflatten_dict(lowerCAmelCase_	)	)
