





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

from utils import logger




class 	UpperCamelCase_							(  UpperCamelCase__     ):


							def __init__(				self     :Any				,					__A     :Any				,					__A     :Any		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   params

									SCREAMING_SNAKE_CASE__								=   np.array(__A		)
									SCREAMING_SNAKE_CASE__								=   np.array([len(__A		) 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     :Dict				,					__A     :Optional[Any]		) -> Optional[Any]:



									"""simple docstring"""



									return (self.token_ids[index], self.lengths[index])


							def __len__(				self     :str		) -> List[Any]:



									"""simple docstring"""



									return len(self.lengths		)


							def   _snake_case	(				self     :Tuple		) -> str:



									"""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   _snake_case	(				self     :List[str]		) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.params.max_model_input_size
									SCREAMING_SNAKE_CASE__								=   self.lengths > max_len
									logger.info(f'''Splitting {sum(__A		)} too long sequences.'''		)

									def divide_chunks(__A     :Dict				,					__A     :Optional[Any]		):
											return [l[i : i + n] for i in range(0				,					len(__A		)				,					__A		)]

									SCREAMING_SNAKE_CASE__								=   []
									SCREAMING_SNAKE_CASE__								=   []
									if self.params.mlm:
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self.params.special_tok_ids["""cls_token"""], self.params.special_tok_ids["""sep_token"""]
									else:
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   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:
													SCREAMING_SNAKE_CASE__								=   []
													for sub_s in divide_chunks(seq_				,					max_len - 2		):
															if sub_s[0] != cls_id:
																	SCREAMING_SNAKE_CASE__								=   np.insert(__A				,					0				,					__A		)
															if sub_s[-1] != sep_id:
																	SCREAMING_SNAKE_CASE__								=   np.insert(__A				,					len(__A		)				,					__A		)
															assert len(__A		) <= max_len
															assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
															sub_seqs.append(__A		)

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

									SCREAMING_SNAKE_CASE__								=   np.array(__A		)
									SCREAMING_SNAKE_CASE__								=   np.array(__A		)


							def   _snake_case	(				self     :str		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   len(self		)
									SCREAMING_SNAKE_CASE__								=   self.lengths > 11
									SCREAMING_SNAKE_CASE__								=   self.token_ids[indices]
									SCREAMING_SNAKE_CASE__								=   self.lengths[indices]
									SCREAMING_SNAKE_CASE__								=   len(self		)
									logger.info(f'''Remove {init_size - new_size} too short (<=11 tokens) sequences.'''		)


							def   _snake_case	(				self     :str		) -> Any:



									"""simple docstring"""



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


							def   _snake_case	(				self     :Union[str, Any]		) -> List[str]:



									"""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   _snake_case	(				self     :List[Any]				,					__A     :Dict		) -> Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   [t[0] for t in batch]
									SCREAMING_SNAKE_CASE__								=   [t[1] for t in batch]
									assert len(__A		) == len(__A		)

									# Max for paddings
									SCREAMING_SNAKE_CASE__								=   max(__A		)

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

									SCREAMING_SNAKE_CASE__								=   torch.tensor(tk_		)  # (bs, max_seq_len_)
									SCREAMING_SNAKE_CASE__								=   torch.tensor(__A		)  # (bs)
									return tk_t, lg_t




'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCamelCase											=		{'configuration_ibert': ['IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'IBertConfig', 'IBertOnnxConfig']}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase											=		[
       'IBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'IBertForMaskedLM',
       'IBertForMultipleChoice',
       'IBertForQuestionAnswering',
       'IBertForSequenceClassification',
       'IBertForTokenClassification',
       'IBertModel',
       'IBertPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_ibert import IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, IBertConfig, IBertOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_ibert import (
          IBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          IBertForMaskedLM,
          IBertForMultipleChoice,
          IBertForQuestionAnswering,
          IBertForSequenceClassification,
          IBertForTokenClassification,
          IBertModel,
          IBertPreTrainedModel,
      )

else:
   import sys

   _UpperCamelCase											=		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






