"""simple docstring"""







import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


if TYPE_CHECKING:
	from transformers.pipelines.conversational import Conversation

_a	:	str=					logging.get_logger(__name__)

_a	:	Tuple=					{
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
}

_a	:	str=					{
    "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",
    },
}

_a	:	List[Any]=					{
    "gpt2": 1_024,
    "gpt2-medium": 1_024,
    "gpt2-large": 1_024,
    "gpt2-xl": 1_024,
    "distilgpt2": 1_024,
}
@lru_cache()
def     __UpperCAmelCase							(       )				->     Optional[Any]:

	'''simple docstring'''





	__snake_case     :				Any 						=					(
	    list(range(ord('!'    )     ,				ord('~'    ) + 1    )    ) + list(range(ord('¡'    )     ,				ord('¬'    ) + 1    )    ) + list(range(ord('®'    )     ,				ord('ÿ'    ) + 1    )    )
	)
	__snake_case     :				List[Any] 						=					bs[:]
	__snake_case     :				List[str] 						=					0
	for b in range(2**8    ):
		if b not in bs:
			bs.append(UpperCAmelCase_    )
			cs.append(2**8 + n    )
			n += 1
	__snake_case     :				str 						=					[chr(UpperCAmelCase_    ) for n in cs]
	return dict(zip(UpperCAmelCase_     ,				UpperCAmelCase_    )    )
def     __UpperCAmelCase							(       UpperCAmelCase_       :							Any    )				->     Dict:

	'''simple docstring'''





	__snake_case     :				List[Any] 						=					set()
	__snake_case     :				Optional[int] 						=					word[0]
	for char in word[1:]:
		pairs.add((prev_char, char)    )
		__snake_case     :				Dict 						=					char
	return pairs

class 						UpperCamelCase     (       lowercase		):
				UpperCAmelCase							: Optional[int]       =						VOCAB_FILES_NAMES
				UpperCAmelCase							: Union[str, Any]       =						PRETRAINED_VOCAB_FILES_MAP
				UpperCAmelCase							: Union[str, Any]       =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				UpperCAmelCase							: Optional[Any]       =						["""input_ids""", """attention_mask"""]



				def __init__(self							: List[Any]      ,    _A							: Any      ,    _A							: Optional[Any]      ,    _A							: Optional[int]="replace"      ,    _A							: Optional[int]="<|endoftext|>"      ,    _A							: Dict="<|endoftext|>"      ,    _A							: Optional[Any]="<|endoftext|>"      ,    _A							: Tuple=None      ,    _A							: List[str]=False      ,    _A							: int=False      ,    **_A							: Optional[Any]      ,    )       ->      Dict:
					__snake_case     :				List[Any] 						=					AddedToken(_A      ,    lstrip=_A      ,    rstrip=_A) if isinstance(_A      ,    _A) else bos_token
					__snake_case     :				Union[str, Any] 						=					AddedToken(_A      ,    lstrip=_A      ,    rstrip=_A) if isinstance(_A      ,    _A) else eos_token
					__snake_case     :				List[Any] 						=					AddedToken(_A      ,    lstrip=_A      ,    rstrip=_A) if isinstance(_A      ,    _A) else unk_token
					__snake_case     :				int 						=					AddedToken(_A      ,    lstrip=_A      ,    rstrip=_A) if isinstance(_A      ,    _A) else pad_token
					super().__init__(
					    errors=_A      ,    unk_token=_A      ,    bos_token=_A      ,    eos_token=_A      ,    pad_token=_A      ,    add_prefix_space=_A      ,    add_bos_token=_A      ,    **_A      ,    )
					__snake_case     :				Tuple 						=					add_bos_token

					with open(_A      ,    encoding='utf-8') as vocab_handle:
						__snake_case     :				str 						=					json.load(_A)
					__snake_case     :				Optional[Any] 						=					{v: k for k, v in self.encoder.items()}
					__snake_case     :				Tuple 						=					errors  # how to handle errors in decoding
					__snake_case     :				Optional[int] 						=					bytes_to_unicode()
					__snake_case     :				Union[str, Any] 						=					{v: k for k, v in self.byte_encoder.items()}
					with open(_A      ,    encoding='utf-8') as merges_handle:
						__snake_case     :				List[Any] 						=					merges_handle.read().split('\n')[1:-1]
					__snake_case     :				List[Any] 						=					[tuple(merge.split()) for merge in bpe_merges]
					__snake_case     :				str 						=					dict(zip(_A      ,    range(len(_A))))
					__snake_case     :				Tuple 						=					{}
					__snake_case     :				List[Any] 						=					add_prefix_space

					# Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
					__snake_case     :				List[Any] 						=					re.compile(r'\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+')



				@property
				def    _lowercase	(self							: Any)       ->      int:
					return len(self.encoder)



				def    _lowercase	(self							: int)       ->      int:
					return dict(self.encoder      ,    **self.added_tokens_encoder)



				def    _lowercase	(self							: Any      ,    _A							: str)       ->      Optional[Any]:
					if token in self.cache:
						return self.cache[token]
					__snake_case     :				Any 						=					tuple(_A)
					__snake_case     :				List[Any] 						=					get_pairs(_A)

					if not pairs:
						return token

					while True:
						__snake_case     :				str 						=					min(_A      ,    key=lambda _A: self.bpe_ranks.get(_A      ,    float('inf')))
						if bigram not in self.bpe_ranks:
							break
						__snake_case						,							__snake_case     :				Union[str, Any] 						=					bigram
						__snake_case     :				Optional[Any] 						=					[]
						__snake_case     :				List[str] 						=					0
						while i < len(_A):
							try:
								__snake_case     :				Optional[Any] 						=					word.index(_A      ,    _A)
							except ValueError:
								new_word.extend(word[i:])
								break
							else:
								new_word.extend(word[i:j])
								__snake_case     :				Optional[Any] 						=					j

							if word[i] == first and i < len(_A) - 1 and word[i + 1] == second:
								new_word.append(first + second)
								i += 2
							else:
								new_word.append(word[i])
								i += 1
						__snake_case     :				List[Any] 						=					tuple(_A)
						__snake_case     :				int 						=					new_word
						if len(_A) == 1:
							break
						else:
							__snake_case     :				List[Any] 						=					get_pairs(_A)
					__snake_case     :				Tuple 						=					' '.join(_A)
					__snake_case     :				str 						=					word
					return word



				def    _lowercase	(self							: List[Any]      ,    _A							: List[str]      ,    _A							: Tuple=None)       ->      Optional[int]:
					if self.add_bos_token:
						__snake_case     :				List[str] 						=					[self.bos_token_id]
					else:
						__snake_case     :				List[str] 						=					[]

					__snake_case     :				List[Any] 						=					bos_token_ids + token_ids_a

					if token_ids_a is None:
						return output

					return output + bos_token_ids + token_ids_a



				def    _lowercase	(self							: str      ,    _A							: List[int]      ,    _A							: Optional[List[int]] = None      ,    _A							: bool = False)       ->      List[int]:
					if already_has_special_tokens:
						return super().get_special_tokens_mask(
						    token_ids_a=_A      ,    token_ids_a=_A      ,    already_has_special_tokens=_A)

					if not self.add_bos_token:
						return super().get_special_tokens_mask(
						    token_ids_a=_A      ,    token_ids_a=_A      ,    already_has_special_tokens=_A)

					if token_ids_a is None:
						return [1] + ([0] * len(_A))
					return [1] + ([0] * len(_A)) + [1] + ([0] * len(_A))



				def    _lowercase	(self							: Dict      ,    _A							: Optional[int])       ->      Tuple:
					__snake_case     :				List[str] 						=					[]
					for token in re.findall(self.pat      ,    _A):
						__snake_case     :				Any 						=					''.join(
						    self.byte_encoder[b] for b in token.encode('utf-8'))  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
						bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_A).split(' '))
					return bpe_tokens



				def    _lowercase	(self							: Tuple      ,    _A							: Dict)       ->      Union[str, Any]:
					return self.encoder.get(_A      ,    self.encoder.get(self.unk_token))



				def    _lowercase	(self							: Optional[int]      ,    _A							: int)       ->      Dict:
					return self.decoder.get(_A)



				def    _lowercase	(self							: Dict      ,    _A							: List[Any])       ->      List[str]:
					__snake_case     :				Any 						=					''.join(_A)
					__snake_case     :				Any 						=					bytearray([self.byte_decoder[c] for c in text]).decode('utf-8'      ,    errors=self.errors)
					return text



				def    _lowercase	(self							: List[Any]      ,    _A							: str      ,    _A							: Optional[str] = None)       ->      Tuple[str]:
					if not os.path.isdir(_A):
						logger.error(f"Vocabulary path ({save_directory}) should be a directory")
						return
					__snake_case     :				Tuple 						=					os.path.join(
					    _A      ,    (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])
					__snake_case     :				int 						=					os.path.join(
					    _A      ,    (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file'])

					with open(_A      ,    'w'      ,    encoding='utf-8') as f:
						f.write(json.dumps(self.encoder      ,    indent=2      ,    sort_keys=_A      ,    ensure_ascii=_A) + '\n')

					__snake_case     :				List[Any] 						=					0
					with open(_A      ,    'w'      ,    encoding='utf-8') as writer:
						writer.write('#version: 0.2\n')
						for bpe_tokens, token_index in sorted(self.bpe_ranks.items()      ,    key=lambda _A: kv[1]):
							if index != token_index:
								logger.warning(
								    f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
								    ' Please check that the tokenizer is not corrupted!')
								__snake_case     :				Optional[Any] 						=					token_index
							writer.write(' '.join(_A) + '\n')
							index += 1

					return vocab_file, merge_file



				def    _lowercase	(self							: Optional[Any]      ,    _A							: Optional[Any]      ,    _A							: Any=False      ,    **_A							: Any)       ->      Any:
					__snake_case     :				str 						=					kwargs.pop('add_prefix_space'      ,    self.add_prefix_space)
					if is_split_into_words or add_prefix_space:
						__snake_case     :				Optional[int] 						=					' ' + text
					return (text, kwargs)






				def    _lowercase	(self							: List[str]      ,    _A							: "Conversation")       ->      List[int]:
					__snake_case     :				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:
						__snake_case     :				str 						=					input_ids[-self.model_max_length :]
					return input_ids






"""simple docstring"""







from typing import TypeVar

from .arrow_dataset import Dataset, _split_by_node_map_style_dataset
from .iterable_dataset import IterableDataset, _split_by_node_iterable_dataset


_a	:	List[Any]=					TypeVar("DatasetType", Dataset, IterableDataset)
def     __UpperCAmelCase							(       UpperCAmelCase_       :							DatasetType     ,				UpperCAmelCase_       :							int     ,				UpperCAmelCase_       :							int    )				->     DatasetType:

	'''simple docstring'''





	if isinstance(UpperCAmelCase_     ,				UpperCAmelCase_    ):
		return _split_by_node_map_style_dataset(UpperCAmelCase_     ,				rank=UpperCAmelCase_     ,				world_size=UpperCAmelCase_    )
	else:
		return _split_by_node_iterable_dataset(UpperCAmelCase_     ,				rank=UpperCAmelCase_     ,				world_size=UpperCAmelCase_    )





