






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

from utils import logger


class     __SCREAMING_SNAKE_CASE					(   __lowercase):






						def __init__(	self     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						params

												lowerCAmelCase__    						=						np.array(_UpperCamelCase   )
												lowerCAmelCase__    						=						np.array([len(_UpperCamelCase   ) for t in data]   )

												self.check()
												self.remove_long_sequences()
												self.remove_empty_sequences()
												self.remove_unknown_sequences()
												self.check()
												self.print_statistics()






						def __getitem__(	self     ,						_UpperCamelCase   ):






												"""simple docstring"""
												return (self.token_ids[index], self.lengths[index])






						def __len__(	self   ):






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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												assert len(self.token_ids   ) == len(self.lengths   )
												assert all(self.lengths[i] == len(self.token_ids[i]   ) for i in range(len(self.lengths   )   )   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.params.max_model_input_size
												lowerCAmelCase__    						=						self.lengths > max_len
												logger.info(F"Splitting {sum(_UpperCamelCase   )} too long sequences."   )

												def divide_chunks(_UpperCamelCase     ,						_UpperCamelCase   ):
																		return [l[i : i + n] for i in range(0     ,						len(_UpperCamelCase   )     ,						_UpperCamelCase   )]

												lowerCAmelCase__    						=						[]
												lowerCAmelCase__    						=						[]
												if self.params.mlm:
																		lowerCAmelCase__						,       lowerCAmelCase__    						=						self.params.special_tok_ids['cls_token'], self.params.special_tok_ids['sep_token']
												else:
																		lowerCAmelCase__						,       lowerCAmelCase__    						=						self.params.special_tok_ids['bos_token'], self.params.special_tok_ids['eos_token']

												for seq_, len_ in zip(self.token_ids     ,						self.lengths   ):
																		assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
																		if len_ <= max_len:
																								new_tok_ids.append(seq_   )
																								new_lengths.append(len_   )
																		else:
																								lowerCAmelCase__    						=						[]
																								for sub_s in divide_chunks(seq_     ,						max_len - 2   ):
																														if sub_s[0] != cls_id:
																																				lowerCAmelCase__    						=						np.insert(_UpperCamelCase     ,						0     ,						_UpperCamelCase   )
																														if sub_s[-1] != sep_id:
																																				lowerCAmelCase__    						=						np.insert(_UpperCamelCase     ,						len(_UpperCamelCase   )     ,						_UpperCamelCase   )
																														assert len(_UpperCamelCase   ) <= max_len
																														assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
																														sub_seqs.append(_UpperCamelCase   )

																								new_tok_ids.extend(_UpperCamelCase   )
																								new_lengths.extend([len(_UpperCamelCase   ) for l in sub_seqs]   )

												lowerCAmelCase__    						=						np.array(_UpperCamelCase   )
												lowerCAmelCase__    						=						np.array(_UpperCamelCase   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						len(self   )
												lowerCAmelCase__    						=						self.lengths > 11
												lowerCAmelCase__    						=						self.token_ids[indices]
												lowerCAmelCase__    						=						self.lengths[indices]
												lowerCAmelCase__    						=						len(self   )
												logger.info(F"Remove {init_size - new_size} too short (<=11 tokens) sequences."   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												if "unk_token" not in self.params.special_tok_ids:
																		return
												else:
																		lowerCAmelCase__    						=						self.params.special_tok_ids['unk_token']
												lowerCAmelCase__    						=						len(self   )
												lowerCAmelCase__    						=						np.array([np.count_nonzero(a == unk_token_id   ) for a in self.token_ids]   )
												lowerCAmelCase__    						=						(unk_occs / self.lengths) < 0.5
												lowerCAmelCase__    						=						self.token_ids[indices]
												lowerCAmelCase__    						=						self.lengths[indices]
												lowerCAmelCase__    						=						len(self   )
												logger.info(F"Remove {init_size - new_size} sequences with a high level of unknown tokens (50%)."   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												if not self.params.is_master:
																		return
												logger.info(F"{len(self   )} sequences"   )
												# data_len = sum(self.lengths)
												# nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
												# logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

												# unk_idx = self.params.special_tok_ids['unk_token']
												# nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
												# logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						[t[0] for t in batch]
												lowerCAmelCase__    						=						[t[1] for t in batch]
												assert len(_UpperCamelCase   ) == len(_UpperCamelCase   )

												# Max for paddings
												lowerCAmelCase__    						=						max(_UpperCamelCase   )

												# Pad token ids
												if self.params.mlm:
																		lowerCAmelCase__    						=						self.params.special_tok_ids['pad_token']
												else:
																		lowerCAmelCase__    						=						self.params.special_tok_ids['unk_token']
												lowerCAmelCase__    						=						[list(t.astype(_UpperCamelCase   )   ) + [pad_idx] * (max_seq_len_ - len(_UpperCamelCase   )) for t in token_ids]
												assert len(tk_   ) == len(_UpperCamelCase   )
												assert all(len(_UpperCamelCase   ) == max_seq_len_ for t in tk_   )

												lowerCAmelCase__    						=						torch.tensor(tk_   )  # (bs, max_seq_len_)
												lowerCAmelCase__    						=						torch.tensor(_UpperCamelCase   )  # (bs)
												return tk_t, lg_t













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


if TYPE_CHECKING:
					from transformers.pipelines.conversational import Conversation


__snake_case    :    Any         =     logging.get_logger(__name__)

__snake_case    :    Tuple         =     {"""tokenizer_file""": """tokenizer.json"""}

__snake_case    :    str         =     {
    """tokenizer_file""": {
        """bigscience/tokenizer""": """https://huggingface.co/bigscience/tokenizer/blob/main/tokenizer.json""",
        """bigscience/bloom-560m""": """https://huggingface.co/bigscience/bloom-560m/blob/main/tokenizer.json""",
        """bigscience/bloom-1b1""": """https://huggingface.co/bigscience/bloom-1b1/blob/main/tokenizer.json""",
        """bigscience/bloom-1b7""": """https://huggingface.co/bigscience/bloom-1b7/blob/main/tokenizer.json""",
        """bigscience/bloom-3b""": """https://huggingface.co/bigscience/bloom-3b/blob/main/tokenizer.json""",
        """bigscience/bloom-7b1""": """https://huggingface.co/bigscience/bloom-7b1/blob/main/tokenizer.json""",
        """bigscience/bloom""": """https://huggingface.co/bigscience/bloom/blob/main/tokenizer.json""",
    },
}


class     __SCREAMING_SNAKE_CASE					(   __lowercase):
						_SCREAMING_SNAKE_CASE	:       Dict   =					VOCAB_FILES_NAMES
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					PRETRAINED_VOCAB_FILES_MAP
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					['''input_ids''', '''attention_mask''']
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					None






						def __init__(	self     ,						_UpperCamelCase=None     ,						_UpperCamelCase=None     ,						_UpperCamelCase=None     ,						_UpperCamelCase="<unk>"     ,						_UpperCamelCase="<s>"     ,						_UpperCamelCase="</s>"     ,						_UpperCamelCase="<pad>"     ,						_UpperCamelCase=False     ,						_UpperCamelCase=False     ,						**_UpperCamelCase     ,						):






												"""simple docstring"""
												super().__init__(
												    _UpperCamelCase     ,						_UpperCamelCase     ,						tokenizer_file=_UpperCamelCase     ,						unk_token=_UpperCamelCase     ,						bos_token=_UpperCamelCase     ,						eos_token=_UpperCamelCase     ,						pad_token=_UpperCamelCase     ,						add_prefix_space=_UpperCamelCase     ,						clean_up_tokenization_spaces=_UpperCamelCase     ,						**_UpperCamelCase     ,						)
												lowerCAmelCase__    						=						json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()   )
												if pre_tok_state.get('add_prefix_space'     ,						_UpperCamelCase   ) != add_prefix_space:
																		lowerCAmelCase__    						=						getattr(_UpperCamelCase     ,						pre_tok_state.pop('type'   )   )
																		lowerCAmelCase__    						=						add_prefix_space
																		lowerCAmelCase__    						=						pre_tok_class(**_UpperCamelCase   )

												lowerCAmelCase__    						=						add_prefix_space






						def       UpperCamelCase__      (	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						kwargs.get('is_split_into_words'     ,						_UpperCamelCase   )
												if not (self.add_prefix_space or not is_split_into_words):
																		raise Exception(
																		    F"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
																		    ' pretokenized inputs.'   )

												return super()._batch_encode_plus(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						kwargs.get('is_split_into_words'     ,						_UpperCamelCase   )

												if not (self.add_prefix_space or not is_split_into_words):
																		raise Exception(
																		    F"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
																		    ' pretokenized inputs.'   )

												return super()._encode_plus(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase = None   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self._tokenizer.model.save(_UpperCamelCase     ,						name=_UpperCamelCase   )
												return tuple(_UpperCamelCase   )







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						[]
												for is_user, text in conversation.iter_texts():
																		input_ids.extend(self.encode(_UpperCamelCase     ,						add_special_tokens=_UpperCamelCase   ) + [self.eos_token_id]   )

												if len(_UpperCamelCase   ) > self.model_max_length:
																		lowerCAmelCase__    						=						input_ids[-self.model_max_length :]
												return input_ids





