




"""simple docstring"""




import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __init__(			self						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   dataset
							lowerCamelCase      =   process
							lowerCamelCase      =   params





				def __len__(			self					):





							"""simple docstring"""
							return len(self.dataset					)





				def __getitem__(			self						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   self.dataset[i]
							lowerCamelCase      =   self.process(_a						,					**self.params					)
							return processed






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __init__(			self						,					_a						,					_a						,					_a						,					_a=None					):





							"""simple docstring"""
							lowerCamelCase      =   loader
							lowerCamelCase      =   infer
							lowerCamelCase      =   params
							if loader_batch_size == 1:
										# Let's spare some time by deactivating altogether
										lowerCamelCase      =   None
							lowerCamelCase      =   loader_batch_size

							# Internal bookkeeping
							lowerCamelCase      =   None
							lowerCamelCase      =   None





				def __len__(			self					):





							"""simple docstring"""
							return len(self.loader					)





				def __iter__(			self					):





							"""simple docstring"""
							lowerCamelCase      =   iter(self.loader					)
							return self





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							if isinstance(self._loader_batch_data						,					torch.Tensor					):
										# Batch data is simple tensor, just fetch the slice
										lowerCamelCase      =   self._loader_batch_data[self._loader_batch_index]
							else:
										# Batch data is assumed to be BaseModelOutput (or dict)
										lowerCamelCase      =   {}
										for k, element in self._loader_batch_data.items():
													if isinstance(_a						,					_a					):
																# Convert ModelOutput to tuple first
																lowerCamelCase      =   element.to_tuple()
																if isinstance(element[0]						,					torch.Tensor					):
																			lowerCamelCase      =   tuple(el[self._loader_batch_index].unsqueeze(0					) for el in element					)
																elif isinstance(element[0]						,					np.ndarray					):
																			lowerCamelCase      =   tuple(np.expand_dims(el[self._loader_batch_index]						,					0					) for el in element					)
																continue
													if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(_a						,					_a					):
																# Those are stored as lists of tensors so need specific unbatching.
																if isinstance(element[0]						,					torch.Tensor					):
																			lowerCamelCase      =   tuple(el[self._loader_batch_index].unsqueeze(0					) for el in element					)
																elif isinstance(element[0]						,					np.ndarray					):
																			lowerCamelCase      =   tuple(np.expand_dims(el[self._loader_batch_index]						,					0					) for el in element					)
																continue
													if element is None:
																# This can happen for optional data that get passed around
																lowerCamelCase      =   None
													elif isinstance(element[self._loader_batch_index]						,					torch.Tensor					):
																# Take correct batch data, but make it looked like batch_size=1
																# For compatibility with other methods within transformers

																lowerCamelCase      =   element[self._loader_batch_index].unsqueeze(0					)
													elif isinstance(element[self._loader_batch_index]						,					np.ndarray					):
																# Take correct batch data, but make it looked like batch_size=1
																# For compatibility with other methods within transformers
																lowerCamelCase      =   np.expand_dims(element[self._loader_batch_index]						,					0					)
													else:
																# This is typically a list, so no need to `unsqueeze`.
																lowerCamelCase      =   element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
										lowerCamelCase      =   self._loader_batch_data.__class__(_a					)
							self._loader_batch_index += 1
							return result





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
										# We are currently unrolling a batch so we just need to return
										# the current item within a batch
										return self.loader_batch_item()

							# We're out of items within a batch
							lowerCamelCase      =   next(self.iterator					)
							lowerCamelCase      =   self.infer(_a						,					**self.params					)
							# We now have a batch of "inferred things".
							if self.loader_batch_size is not None:
										# Try to infer the size of the batch
										if isinstance(_a						,					torch.Tensor					):
													lowerCamelCase      =   processed
										else:
													lowerCamelCase      =   list(processed.keys()					)[0]
													lowerCamelCase      =   processed[key]
										if isinstance(_a						,					_a					):
													lowerCamelCase      =   len(_a					)
										else:
													lowerCamelCase      =   first_tensor.shape[0]
										if 0 < observed_batch_size < self.loader_batch_size:
													# could be last batch so we can't unroll as many
													# elements.
													lowerCamelCase      =   observed_batch_size
										# Setting internal index to unwrap the batch
										lowerCamelCase      =   processed
										lowerCamelCase      =   0
										return self.loader_batch_item()
							else:
										# We're not unrolling batches
										return processed






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __init__(			self						,					_a						,					_a						,					_a						,					_a=None					):





							"""simple docstring"""
							super().__init__(_a						,					_a						,					_a					)





				def __iter__(			self					):





							"""simple docstring"""
							lowerCamelCase      =   iter(self.loader					)
							lowerCamelCase      =   None
							return self





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							if self.subiterator is None:
										lowerCamelCase      =   self.infer(next(self.iterator					)						,					**self.params					)
							try:
										# Try to return next item
										lowerCamelCase      =   next(self.subiterator					)
							except StopIteration:
										# When a preprocess iterator ends, we can start lookig at the next item
										# ChunkIterator will keep feeding until ALL elements of iterator
										# all have created their subiterator and have been iterating against.
										#
										# Another way to look at it, is we're basically flattening lists of lists
										# into a single list, but with generators
										lowerCamelCase      =   self.infer(next(self.iterator					)						,					**self.params					)
										lowerCamelCase      =   next(self.subiterator					)
							return processed






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __iter__(			self					):





							"""simple docstring"""
							lowerCamelCase      =   iter(self.loader					)
							return self





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# Extremely similar to PipelineIterator in its unpacking mechanism
							# BUT, we have an extra required item which is the presence of `is_last`
							# That is because everything is flattened by `PipelineChunkIterator` we
							# need to keep track of how to regroup here in the original `process`
							# boundaries so that `process` and `postprocess` see the same data.

							# This iterator accumulates items (possibly while unbatching) until it
							# its a `is_last` and then just passes it on to the caller.
							lowerCamelCase      =   False
							lowerCamelCase      =   []
							if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
										while self._loader_batch_index < self.loader_batch_size:
													lowerCamelCase      =   self.loader_batch_item()
													lowerCamelCase      =   item.pop("""is_last"""					)
													accumulator.append(_a					)
													if is_last:
																return accumulator

							while not is_last:
										lowerCamelCase      =   self.infer(next(self.iterator					)						,					**self.params					)
										if self.loader_batch_size is not None:
													if isinstance(_a						,					torch.Tensor					):
																lowerCamelCase      =   processed
													else:
																lowerCamelCase      =   list(processed.keys()					)[0]
																lowerCamelCase      =   processed[key]
													if isinstance(_a						,					_a					):
																lowerCamelCase      =   len(_a					)
													else:
																lowerCamelCase      =   first_tensor.shape[0]
													if 0 < observed_batch_size < self.loader_batch_size:
																# could be last batch so we can't unroll as many
																# elements.
																lowerCamelCase      =   observed_batch_size
													lowerCamelCase      =   processed
													lowerCamelCase      =   0
													while self._loader_batch_index < self.loader_batch_size:
																lowerCamelCase      =   self.loader_batch_item()
																lowerCamelCase      =   item.pop("""is_last"""					)
																accumulator.append(_a					)
																if is_last:
																			return accumulator
										else:
													lowerCamelCase      =   processed
													lowerCamelCase      =   item.pop("""is_last"""					)
													accumulator.append(_a					)
							return accumulator






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __init__(			self						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   dataset
							lowerCamelCase      =   key





				def __len__(			self					):





							"""simple docstring"""
							return len(self.dataset					)





				def __getitem__(			self						,					_a					):





							"""simple docstring"""
							return self.dataset[i][self.key]






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __init__(			self						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   dataset
							lowerCamelCase      =   keya
							lowerCamelCase      =   keya





				def __len__(			self					):





							"""simple docstring"""
							return len(self.dataset					)





				def __getitem__(			self						,					_a					):





							"""simple docstring"""
							return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}









"""simple docstring"""




import json
import os
from typing import Optional, Tuple

import regex as re

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCAmelCase				:     Union[str, Any]					=				logging.get_logger(__name__)

lowerCAmelCase				:     Optional[int]					=				{
    """vocab_file""": """vocab.json""",
    """merges_file""": """merges.txt""",
}

lowerCAmelCase				:     Optional[Any]					=				{
    """vocab_file""": {"""ctrl""": """https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-vocab.json"""},
    """merges_file""": {"""ctrl""": """https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-merges.txt"""},
}

lowerCAmelCase				:     Optional[Any]					=				{
    """ctrl""": 256,
}

lowerCAmelCase				:     int					=				{
    """Pregnancy""": 16_8629,
    """Christianity""": 7675,
    """Explain""": 10_6423,
    """Fitness""": 6_3440,
    """Saving""": 6_3163,
    """Ask""": 2_7171,
    """Ass""": 9_5985,
    """Joke""": 16_3509,
    """Questions""": 4_5622,
    """Thoughts""": 4_9605,
    """Retail""": 5_2342,
    """Feminism""": 16_4338,
    """Writing""": 1_1992,
    """Atheism""": 19_2263,
    """Netflix""": 4_8616,
    """Computing""": 3_9639,
    """Opinion""": 4_3213,
    """Alone""": 4_4967,
    """Funny""": 5_8917,
    """Gaming""": 4_0358,
    """Human""": 4088,
    """India""": 1331,
    """Joker""": 7_7138,
    """Diet""": 3_6206,
    """Legal""": 1_1859,
    """Norman""": 4939,
    """Tip""": 7_2689,
    """Weight""": 5_2343,
    """Movies""": 4_6273,
    """Running""": 2_3425,
    """Science""": 2090,
    """Horror""": 3_7793,
    """Confession""": 6_0572,
    """Finance""": 1_2250,
    """Politics""": 1_6360,
    """Scary""": 19_1985,
    """Support""": 1_2654,
    """Technologies""": 3_2516,
    """Teenage""": 6_6160,
    """Event""": 3_2769,
    """Learned""": 6_7460,
    """Notion""": 18_2770,
    """Wikipedia""": 3_7583,
    """Books""": 6665,
    """Extract""": 7_6050,
    """Confessions""": 10_2701,
    """Conspiracy""": 7_5932,
    """Links""": 6_3674,
    """Narcissus""": 15_0425,
    """Relationship""": 5_4766,
    """Relationships""": 13_4796,
    """Reviews""": 4_1671,
    """News""": 4256,
    """Translation""": 2_6820,
    """multilingual""": 12_8406,
}



def        a__					(							snake_case__							)		->   Optional[Any]:
			lowerCamelCase      =   set()
			lowerCamelCase      =   word[0]
			for char in word[1:]:
						pairs.add((prev_char, char)							)
						lowerCamelCase      =   char

			lowerCamelCase      =   set(snake_case__							)
			return pairs






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''
				__UpperCamelCase         =		VOCAB_FILES_NAMES
				__UpperCamelCase         =		PRETRAINED_VOCAB_FILES_MAP
				__UpperCamelCase         =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				__UpperCamelCase         =		CONTROL_CODES





				def __init__(			self						,					_a						,					_a						,					_a="<unk>"						,					**_a					):





							"""simple docstring"""
							super().__init__(unk_token=_a						,					**_a					)

							with open(_a						,					encoding="""utf-8"""					) as vocab_handle:
										lowerCamelCase      =   json.load(_a					)
							lowerCamelCase      =   {v: k for k, v in self.encoder.items()}
							with open(_a						,					encoding="""utf-8"""					) as merges_handle:
										lowerCamelCase      =   merges_handle.read().split("""\n"""					)[1:-1]
							lowerCamelCase      =   [tuple(merge.split()					) for merge in merges]
							lowerCamelCase      =   dict(zip(_a						,					range(len(_a					)					)					)					)
							lowerCamelCase      =   {}





				@property
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return len(self.encoder					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return dict(self.encoder						,					**self.added_tokens_encoder					)





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							if token in self.cache:
										return self.cache[token]
							lowerCamelCase      =   tuple(_a					)
							lowerCamelCase      =   tuple(list(word[:-1]					) + [word[-1] + """</w>"""]					)
							lowerCamelCase      =   get_pairs(_a					)

							if not pairs:
										return token

							while True:
										lowerCamelCase      =   min(_a						,					key=lambda _a					: self.bpe_ranks.get(_a						,					float("""inf"""					)					)					)
										if bigram not in self.bpe_ranks:
													break
										lowerCamelCase ,  lowerCamelCase      =   bigram
										lowerCamelCase      =   []
										lowerCamelCase      =   0
										while i < len(_a					):
													try:
																lowerCamelCase      =   word.index(_a						,					_a					)
													except ValueError:
																new_word.extend(word[i:]					)
																break
													else:
																new_word.extend(word[i:j]					)
																lowerCamelCase      =   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
										lowerCamelCase      =   tuple(_a					)
										lowerCamelCase      =   new_word
										if len(_a					) == 1:
													break
										else:
													lowerCamelCase      =   get_pairs(_a					)
							lowerCamelCase      =   """@@ """.join(_a					)
							lowerCamelCase      =   word[:-4]
							lowerCamelCase      =   word
							return word





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   []

							lowerCamelCase      =   re.findall(r"""\S+\n?"""						,					_a					)

							for token in words:
										split_tokens.extend(list(self.bpe(_a					).split(""" """					)					)					)
							return split_tokens





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							return self.encoder.get(_a						,					self.encoder.get(self.unk_token					)					)





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							return self.decoder.get(_a						,					self.unk_token					)





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   """ """.join(_a					).replace("""@@ """						,					""""""					).strip()
							return out_string





				def   _lowerCAmelCase							(			self						,					_a						,					_a = None					):





							"""simple docstring"""
							if not os.path.isdir(_a					):
										logger.error(f'Vocabulary path ({save_directory}) should be a directory'					)
										return
							lowerCamelCase      =   os.path.join(
							    _a						,					(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]					)
							lowerCamelCase      =   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"""					)

							lowerCamelCase      =   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!"""					)
																lowerCamelCase      =   token_index
													writer.write(""" """.join(_a					) + """\n"""					)
													index += 1

							return vocab_file, merge_file

				# def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
				#     filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
				#     tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
				#     tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
				#     return ''.join(tokens_generated_so_far)



