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

from utils import logger



class    a_ (			_a    ):






			def __init__(      self				,			__UpperCamelCase				,			__UpperCamelCase  ):
									_lowercase					       =					params

									_lowercase					       =					np.array(__UpperCamelCase  )
									_lowercase					       =					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  ):
									return (self.token_ids[index], self.lengths[index])






			def __len__(      self  ):
									return len(self.lengths  )






			def    UpperCamelCase_				(      self  ):
									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  ):
									_lowercase					       =					self.params.max_model_input_size
									_lowercase					       =					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  )]

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

									_lowercase					       =					np.array(__UpperCamelCase  )
									_lowercase					       =					np.array(__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					len(self  )
									_lowercase					       =					self.lengths > 11
									_lowercase					       =					self.token_ids[indices]
									_lowercase					       =					self.lengths[indices]
									_lowercase					       =					len(self  )
									logger.info(f"""Remove {init_size - new_size} too short (<=11 tokens) sequences."""  )






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






			def    UpperCamelCase_				(      self  ):
									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  ):
									_lowercase					       =					[t[0] for t in batch]
									_lowercase					       =					[t[1] for t in batch]
									assert len(__UpperCamelCase  ) == len(__UpperCamelCase  )

									# Max for paddings
									_lowercase					       =					max(__UpperCamelCase  )

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

									_lowercase					       =					torch.tensor(tk_  )  # (bs, max_seq_len_)
									_lowercase					       =					torch.tensor(__UpperCamelCase  )  # (bs)
									return tk_t, lg_t

import datasets


A		: Optional[int]								=     '''\
@InProceedings{conneau2018xnli,
  author = "Conneau, Alexis
                 and Rinott, Ruty
                 and Lample, Guillaume
                 and Williams, Adina
                 and Bowman, Samuel R.
                 and Schwenk, Holger
                 and Stoyanov, Veselin",
  title = "XNLI: Evaluating Cross-lingual Sentence Representations",
  booktitle = "Proceedings of the 2018 Conference on Empirical Methods
               in Natural Language Processing",
  year = "2018",
  publisher = "Association for Computational Linguistics",
  location = "Brussels, Belgium",
}
'''

A		: Optional[int]								=     '''\
XNLI is a subset of a few thousand examples from MNLI which has been translated
into a 14 different languages (some low-ish resource). As with MNLI, the goal is
to predict textual entailment (does sentence A imply/contradict/neither sentence
B) and is a classification task (given two sentences, predict one of three
labels).
'''

A		: str								=     '''
Computes XNLI score which is just simple accuracy.
Args:
    predictions: Predicted labels.
    references: Ground truth labels.
Returns:
    \'accuracy\': accuracy
Examples:

    >>> predictions = [0, 1]
    >>> references = [0, 1]
    >>> xnli_metric = datasets.load_metric("xnli")
    >>> results = xnli_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}
'''




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	List[str]		,   SCREAMING_SNAKE_CASE_	:	List[Any]       )  ->     str:
						return (preds == labels).mean()



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION					,			_KWARGS_DESCRIPTION    )
class    a_ (			datasets.Metric    ):






			def    UpperCamelCase_				(      self  ):
									return datasets.MetricInfo(
									    description=_DESCRIPTION				,			citation=_CITATION				,			inputs_description=_KWARGS_DESCRIPTION				,			features=datasets.Features(
									        {
									            """predictions""": datasets.Value("""int64""" if self.config_name != """sts-b""" else """float32"""  ),
									            """references""": datasets.Value("""int64""" if self.config_name != """sts-b""" else """float32"""  ),
									        }  )				,			codebase_urls=[]				,			reference_urls=[]				,			format="""numpy"""				,			)






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase  ):
									return {"accuracy": simple_accuracy(__UpperCamelCase				,			__UpperCamelCase  )}