



"""simple docstring"""





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

from utils import logger



class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__		):





  """simple docstring"""






  def __init__(					self							,    _lowercase							,    _lowercase   ) ->						Tuple:







        '''simple docstring'''

        snake_case_   : Tuple						=			params

        snake_case_   : Optional[Any]						=			np.array(_lowerCAmelCase   )
        snake_case_   : List[Any]						=			np.array([len(_lowerCAmelCase   ) 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							,    _lowercase   ) ->						List[str]:







        '''simple docstring'''

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





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







        '''simple docstring'''

        return len(self.lengths   )





  def    UpperCAmelCase__							(					self   ) ->						Optional[int]:







        '''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   ) ->						List[str]:







        '''simple docstring'''

        snake_case_   : Any						=			self.params.max_model_input_size
        snake_case_   : Optional[int]						=			self.lengths > max_len
        logger.info(f'Splitting {sum(_lowerCAmelCase   )} too long sequences.'   )

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

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

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

        snake_case_   : Optional[Any]						=			np.array(_lowerCAmelCase   )
        snake_case_   : List[str]						=			np.array(_lowerCAmelCase   )





  def    UpperCAmelCase__							(					self   ) ->						Any:







        '''simple docstring'''

        snake_case_   : Optional[Any]						=			len(self   )
        snake_case_   : Dict						=			self.lengths > 1_1
        snake_case_   : Optional[Any]						=			self.token_ids[indices]
        snake_case_   : List[str]						=			self.lengths[indices]
        snake_case_   : Optional[Any]						=			len(self   )
        logger.info(f'Remove {init_size - new_size} too short (<=11 tokens) sequences.'   )





  def    UpperCAmelCase__							(					self   ) ->						Optional[int]:







        '''simple docstring'''

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





  def    UpperCAmelCase__							(					self   ) ->						int:







        '''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							,    _lowercase   ) ->						List[Any]:







        '''simple docstring'''

        snake_case_   : str						=			[t[0] for t in batch]
        snake_case_   : str						=			[t[1] for t in batch]
        assert len(_lowerCAmelCase   ) == len(_lowerCAmelCase   )

        # Max for paddings
        snake_case_   : Union[str, Any]						=			max(_lowerCAmelCase   )

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

        snake_case_   : str						=			torch.tensor(tk_   )  # (bs, max_seq_len_)
        snake_case_   : Optional[Any]						=			torch.tensor(_lowerCAmelCase   )  # (bs)
        return tk_t, lg_t






"""simple docstring"""





from .data_collator import (
    DataCollatorForLanguageModeling,
    DataCollatorForPermutationLanguageModeling,
    DataCollatorForSeqaSeq,
    DataCollatorForSOP,
    DataCollatorForTokenClassification,
    DataCollatorForWholeWordMask,
    DataCollatorWithPadding,
    DefaultDataCollator,
    default_data_collator,
)
from .metrics import glue_compute_metrics, xnli_compute_metrics
from .processors import (
    DataProcessor,
    InputExample,
    InputFeatures,
    SingleSentenceClassificationProcessor,
    SquadExample,
    SquadFeatures,
    SquadVaProcessor,
    SquadVaProcessor,
    glue_convert_examples_to_features,
    glue_output_modes,
    glue_processors,
    glue_tasks_num_labels,
    squad_convert_examples_to_features,
    xnli_output_modes,
    xnli_processors,
    xnli_tasks_num_labels,
)

