





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

from utils import logger




class        _a    (    SCREAMING_SNAKE_CASE_):

      """simple docstring"""


      def __init__(		self:			int	,    __lowerCamelCase:			str	,    __lowerCamelCase:			Tuple					):
          '''simple docstring'''


          UpperCamelCase__:	Optional[Any]					=     params

          UpperCamelCase__:	Any					=     np.array(_a					)
          UpperCamelCase__:	int					=     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:			str	,    __lowerCamelCase:			Optional[int]					):
          '''simple docstring'''


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






      def __len__(		self:			int					):
          '''simple docstring'''


          return len(self.lengths					)






      def 			UpperCAmelCase_       (		self:			Tuple					):
          '''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:			Optional[int]					):
          '''simple docstring'''


          UpperCamelCase__:	List[str]					=     self.params.max_model_input_size
          UpperCamelCase__:	List[Any]					=     self.lengths > max_len
          logger.info(F"Splitting {sum(_a					)} too long sequences."					)

          def divide_chunks(__lowerCamelCase:			Union[str, Any]	,    __lowerCamelCase:			Any					):
              return [l[i : i + n] for i in range(0	,    len(_a					)	,    _a					)]

          UpperCamelCase__:	Dict					=     []
          UpperCamelCase__:	Tuple					=     []
          if self.params.mlm:
              UpperCamelCase__:	Any					=     self.params.special_tok_ids['cls_token'], self.params.special_tok_ids['sep_token']
          else:
              UpperCamelCase__:	Dict					=     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:
                  UpperCamelCase__:	Tuple					=     []
                  for sub_s in divide_chunks(seq_	,    max_len - 2					):
                      if sub_s[0] != cls_id:
                          UpperCamelCase__:	str					=     np.insert(_a	,    0	,    _a					)
                      if sub_s[-1] != sep_id:
                          UpperCamelCase__:	List[str]					=     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]					)

          UpperCamelCase__:	str					=     np.array(_a					)
          UpperCamelCase__:	int					=     np.array(_a					)






      def 			UpperCAmelCase_       (		self:			List[str]					):
          '''simple docstring'''


          UpperCamelCase__:	Tuple					=     len(self					)
          UpperCamelCase__:	Union[str, Any]					=     self.lengths > 11
          UpperCamelCase__:	Optional[Any]					=     self.token_ids[indices]
          UpperCamelCase__:	Union[str, Any]					=     self.lengths[indices]
          UpperCamelCase__:	List[Any]					=     len(self					)
          logger.info(F"Remove {init_size - new_size} too short (<=11 tokens) sequences."					)






      def 			UpperCAmelCase_       (		self:			str					):
          '''simple docstring'''


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






      def 			UpperCAmelCase_       (		self:			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 			UpperCAmelCase_       (		self:			Union[str, Any]	,    __lowerCamelCase:			int					):
          '''simple docstring'''


          UpperCamelCase__:	Optional[Any]					=     [t[0] for t in batch]
          UpperCamelCase__:	Dict					=     [t[1] for t in batch]
          assert len(_a					) == len(_a					)

          # Max for paddings
          UpperCamelCase__:	Dict					=     max(_a					)

          # Pad token ids
          if self.params.mlm:
              UpperCamelCase__:	int					=     self.params.special_tok_ids['pad_token']
          else:
              UpperCamelCase__:	Optional[int]					=     self.params.special_tok_ids['unk_token']
          UpperCamelCase__:	int					=     [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_					)

          UpperCamelCase__:	Any					=     torch.tensor(tk_					)  # (bs, max_seq_len_)
          UpperCamelCase__:	List[Any]					=     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__)






