




import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
	    DistilBertForMaskedLM,
	    DistilBertForMultipleChoice,
	    DistilBertForQuestionAnswering,
	    DistilBertForSequenceClassification,
	    DistilBertForTokenClassification,
	    DistilBertModel,
	)
class        __snake_case  (_a		):



				def __init__(      self      :		int		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Tuple=13		,	_UpperCAmelCase      :		Optional[Any]=7		,	_UpperCAmelCase      :		Optional[Any]=True		,	_UpperCAmelCase      :		List[Any]=True		,	_UpperCAmelCase      :		Any=False		,	_UpperCAmelCase      :		Tuple=True		,	_UpperCAmelCase      :		str=99		,	_UpperCAmelCase      :		List[Any]=32		,	_UpperCAmelCase      :		Dict=5		,	_UpperCAmelCase      :		List[str]=4		,	_UpperCAmelCase      :		List[Any]=37		,	_UpperCAmelCase      :		List[Any]="gelu"		,	_UpperCAmelCase      :		Any=0.1		,	_UpperCAmelCase      :		Tuple=0.1		,	_UpperCAmelCase      :		str=512		,	_UpperCAmelCase      :		Any=16		,	_UpperCAmelCase      :		Optional[Any]=2		,	_UpperCAmelCase      :		Optional[Any]=0.02		,	_UpperCAmelCase      :		str=3		,	_UpperCAmelCase      :		List[Any]=4		,	_UpperCAmelCase      :		int=None		,	)  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      parent
								_lowerCAmelCase    :    List[str]										=      batch_size
								_lowerCAmelCase    :    str										=      seq_length
								_lowerCAmelCase    :    Optional[Any]										=      is_training
								_lowerCAmelCase    :    int										=      use_input_mask
								_lowerCAmelCase    :    Union[str, Any]										=      use_token_type_ids
								_lowerCAmelCase    :    Optional[int]										=      use_labels
								_lowerCAmelCase    :    Any										=      vocab_size
								_lowerCAmelCase    :    List[str]										=      hidden_size
								_lowerCAmelCase    :    Tuple										=      num_hidden_layers
								_lowerCAmelCase    :    Tuple										=      num_attention_heads
								_lowerCAmelCase    :    Tuple										=      intermediate_size
								_lowerCAmelCase    :    int										=      hidden_act
								_lowerCAmelCase    :    Dict										=      hidden_dropout_prob
								_lowerCAmelCase    :    Tuple										=      attention_probs_dropout_prob
								_lowerCAmelCase    :    List[str]										=      max_position_embeddings
								_lowerCAmelCase    :    Optional[int]										=      type_vocab_size
								_lowerCAmelCase    :    int										=      type_sequence_label_size
								_lowerCAmelCase    :    List[Any]										=      initializer_range
								_lowerCAmelCase    :    Dict										=      num_labels
								_lowerCAmelCase    :    Dict										=      num_choices
								_lowerCAmelCase    :    Optional[int]										=      scope



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    Dict:





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size      )

								_lowerCAmelCase    :    int										=      None
								if self.use_input_mask:
												_lowerCAmelCase    :    str										=      random_attention_mask([self.batch_size, self.seq_length]      )

								_lowerCAmelCase    :    Tuple										=      None
								_lowerCAmelCase    :    List[Any]										=      None
								_lowerCAmelCase    :    Optional[Any]										=      None
								if self.use_labels:
												_lowerCAmelCase    :    List[Any]										=      ids_tensor([self.batch_size]		,	self.type_sequence_label_size      )
												_lowerCAmelCase    :    Optional[int]										=      ids_tensor([self.batch_size, self.seq_length]		,	self.num_labels      )
												_lowerCAmelCase    :    Tuple										=      ids_tensor([self.batch_size]		,	self.num_choices      )

								_lowerCAmelCase    :    Optional[int]										=      self.get_config()

								return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



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





								'''simple docstring'''




								return DistilBertConfig(
								    vocab_size=self.vocab_size		,	dim=self.hidden_size		,	n_layers=self.num_hidden_layers		,	n_heads=self.num_attention_heads		,	hidden_dim=self.intermediate_size		,	hidden_act=self.hidden_act		,	dropout=self.hidden_dropout_prob		,	attention_dropout=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	initializer_range=self.initializer_range		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		str		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Optional[int]		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		Union[str, Any]      )  ->    List[str]:





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      DistilBertModel(config=_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    Dict										=      model(_UpperCAmelCase		,	_UpperCAmelCase      )
								_lowerCAmelCase    :    List[str]										=      model(_UpperCAmelCase      )
								self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, self.seq_length, self.hidden_size)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		Optional[int]		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Optional[Any]      )  ->    Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      DistilBertForMaskedLM(config=_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    str										=      model(_UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	labels=_UpperCAmelCase      )
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.vocab_size)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		Any      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      DistilBertForQuestionAnswering(config=_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    List[Any]										=      model(
								    _UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	start_positions=_UpperCAmelCase		,	end_positions=_UpperCAmelCase      )
								self.parent.assertEqual(result.start_logits.shape		,	(self.batch_size, self.seq_length)      )
								self.parent.assertEqual(result.end_logits.shape		,	(self.batch_size, self.seq_length)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[int]		,	_UpperCAmelCase      :		Optional[Any]		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		Dict      )  ->    List[str]:





								'''simple docstring'''




								_lowerCAmelCase    :    List[Any]										=      self.num_labels
								_lowerCAmelCase    :    str										=      DistilBertForSequenceClassification(_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    Tuple										=      model(_UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	labels=_UpperCAmelCase      )
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.num_labels)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		List[Any]      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      self.num_labels
								_lowerCAmelCase    :    Optional[int]										=      DistilBertForTokenClassification(config=_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()

								_lowerCAmelCase    :    int										=      model(_UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	labels=_UpperCAmelCase      )
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.num_labels)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Union[str, Any]      )  ->    Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      self.num_choices
								_lowerCAmelCase    :    Dict										=      DistilBertForMultipleChoice(config=_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    Tuple										=      input_ids.unsqueeze(1      ).expand(-1		,	self.num_choices		,	-1      ).contiguous()
								_lowerCAmelCase    :    str										=      input_mask.unsqueeze(1      ).expand(-1		,	self.num_choices		,	-1      ).contiguous()
								_lowerCAmelCase    :    int										=      model(
								    _UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	labels=_UpperCAmelCase		,	)
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.num_choices)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      self.prepare_config_and_inputs()
								((_lowerCAmelCase)					,    (_lowerCAmelCase)					,    (_lowerCAmelCase)					,    (_lowerCAmelCase)					,    (_lowerCAmelCase)					,    (_lowerCAmelCase))    :    Dict										=      config_and_inputs
								_lowerCAmelCase    :    int										=      {"""input_ids""": input_ids, """attention_mask""": input_mask}
								return config, inputs_dict






@require_torch
class        __snake_case  (_a						,	_a						,	unittest.TestCase		):
				lowerCAmelCase__				=     (
				    (
				        DistilBertModel,
				        DistilBertForMaskedLM,
				        DistilBertForMultipleChoice,
				        DistilBertForQuestionAnswering,
				        DistilBertForSequenceClassification,
				        DistilBertForTokenClassification,
				    )
				    if is_torch_available()
				    else None
				)
				lowerCAmelCase__				=     (
				    {
				        "feature-extraction": DistilBertModel,
				        "fill-mask": DistilBertForMaskedLM,
				        "question-answering": DistilBertForQuestionAnswering,
				        "text-classification": DistilBertForSequenceClassification,
				        "token-classification": DistilBertForTokenClassification,
				        "zero-shot": DistilBertForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				lowerCAmelCase__				=     True
				lowerCAmelCase__				=     True
				lowerCAmelCase__				=     True
				lowerCAmelCase__				=     True



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]      )  ->    Union[str, Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      DistilBertModelTester(self      )
								_lowerCAmelCase    :    Any										=      ConfigTester(self		,	config_class=_UpperCAmelCase		,	dim=37      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Any:





								'''simple docstring'''




								self.config_tester.run_common_tests()



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





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_model(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    List[Any]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_masked_lm(*_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_question_answering(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_sequence_classification(*_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_token_classification(*_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_multiple_choice(*_UpperCAmelCase      )



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]      )  ->    str:





								'''simple docstring'''




								for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_lowerCAmelCase    :    List[Any]										=      DistilBertModel.from_pretrained(_UpperCAmelCase      )
												self.assertIsNotNone(_UpperCAmelCase      )



				@slow
				@require_torch_gpu
				def       SCREAMING_SNAKE_CASE       (      self      :		Union[str, Any]      )  ->    List[str]:





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    Optional[Any]										=      self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
												# BertForMultipleChoice behaves incorrectly in JIT environments.
												if model_class == DistilBertForMultipleChoice:
																return

												_lowerCAmelCase    :    Optional[int]										=      True
												_lowerCAmelCase    :    List[str]										=      model_class(config=_UpperCAmelCase      )

												_lowerCAmelCase    :    List[str]										=      self._prepare_for_class(_UpperCAmelCase		,	_UpperCAmelCase      )
												_lowerCAmelCase    :    List[Any]										=      torch.jit.trace(
												    _UpperCAmelCase		,	(inputs_dict["""input_ids"""].to("""cpu"""      ), inputs_dict["""attention_mask"""].to("""cpu"""      ))      )

												with tempfile.TemporaryDirectory() as tmp:
																torch.jit.save(_UpperCAmelCase		,	os.path.join(_UpperCAmelCase		,	"""traced_model.pt"""      )      )
																_lowerCAmelCase    :    str										=      torch.jit.load(os.path.join(_UpperCAmelCase		,	"""traced_model.pt"""      )		,	map_location=_UpperCAmelCase      )
																loaded(inputs_dict["""input_ids"""].to(_UpperCAmelCase      )		,	inputs_dict["""attention_mask"""].to(_UpperCAmelCase      )      )






@require_torch
class        __snake_case  (unittest.TestCase		):



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      DistilBertModel.from_pretrained("""distilbert-base-uncased"""      )
								_lowerCAmelCase    :    Optional[int]										=      torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]      )
								_lowerCAmelCase    :    Union[str, Any]										=      torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
								with torch.no_grad():
												_lowerCAmelCase    :    Tuple										=      model(_UpperCAmelCase		,	attention_mask=_UpperCAmelCase      )[0]
								_lowerCAmelCase    :    str										=      torch.Size((1, 11, 768)      )
								self.assertEqual(output.shape		,	_UpperCAmelCase      )
								_lowerCAmelCase    :    int										=      torch.tensor(
								    [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]      )

								self.assertTrue(torch.allclose(output[:, 1:4, 1:4]		,	_UpperCAmelCase		,	atol=1E-4      )      )







import comet  # From: unbabel-comet
import torch

import datasets


_lowerCamelCase      :  List[Any]      						=						datasets.logging.get_logger(__name__)

_lowerCamelCase      :  Optional[Any]      						=						"\\n@inproceedings{rei-EtAl:2020:WMT,\n   author    = {Rei, Ricardo  and  Stewart, Craig  and  Farinha, Ana C  and  Lavie, Alon},\n   title     = {Unbabel's Participation in the WMT20 Metrics Shared Task},\n   booktitle      = {Proceedings of the Fifth Conference on Machine Translation},\n   month          = {November},\n   year           = {2020},\n   address        = {Online},\n   publisher      = {Association for Computational Linguistics},\n   pages     = {909--918},\n}\n@inproceedings{rei-etal-2020-comet,\n   title = \"{COMET}: A Neural Framework for {MT} Evaluation\",\n   author = \"Rei, Ricardo  and\n      Stewart, Craig  and\n      Farinha, Ana C  and\n      Lavie, Alon\",\n   booktitle = \"Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)\",\n   month = nov,\n   year = \"2020\",\n   address = \"Online\",\n   publisher = \"Association for Computational Linguistics\",\n   url = \"https://www.aclweb.org/anthology/2020.emnlp-main.213\",\n   pages = \"2685--2702\",\n}\n"

_lowerCamelCase      :  Union[str, Any]      						=						"\\nCrosslingual Optimized Metric for Evaluation of Translation (COMET) is an open-source framework used to train Machine Translation metrics that achieve high levels of correlation with different types of human judgments (HTER, DA's or MQM).\nWith the release of the framework the authors also released fully trained models that were used to compete in the WMT20 Metrics Shared Task achieving SOTA in that years competition.\n\nSee the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.\n"

_lowerCamelCase      :  Optional[int]      						=						"\nCOMET score.\n\nArgs:\n\n`sources` (list of str): Source sentences\n`predictions` (list of str): candidate translations\n`references` (list of str): reference translations\n`cuda` (bool): If set to True, runs COMET using GPU\n`show_progress` (bool): Shows progress\n`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.\n\nReturns:\n    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.\n    `scores`: List of scores.\n\nExamples:\n\n    >>> comet_metric = datasets.load_metric('comet')\n    >>> # comet_metric = load_metric('comet', 'wmt20-comet-da')  # you can also choose which model to use\n    >>> source = [\"Dem Feuer konnte Einhalt geboten werden\", \"Schulen und Kindergärten wurden eröffnet.\"]\n    >>> hypothesis = [\"The fire could be stopped\", \"Schools and kindergartens were open\"]\n    >>> reference = [\"They were able to control the fire.\", \"Schools and kindergartens opened\"]\n    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)\n    >>> print([round(v, 2) for v in results[\"scores\"]])\n    [0.19, 0.92]\n"
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,	_KWARGS_DESCRIPTION		)
class        __snake_case  (datasets.Metric		):



				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    Tuple:





								'''simple docstring'''




								return datasets.MetricInfo(
								    description=_DESCRIPTION		,	citation=_CITATION		,	homepage="""https://unbabel.github.io/COMET/html/index.html"""		,	inputs_description=_KWARGS_DESCRIPTION		,	features=datasets.Features(
								        {
								            """sources""": datasets.Value("""string"""		,	id="""sequence"""      ),
								            """predictions""": datasets.Value("""string"""		,	id="""sequence"""      ),
								            """references""": datasets.Value("""string"""		,	id="""sequence"""      ),
								        }      )		,	codebase_urls=["""https://github.com/Unbabel/COMET"""]		,	reference_urls=[
								        """https://github.com/Unbabel/COMET""",
								        """https://www.aclweb.org/anthology/2020.emnlp-main.213/""",
								        """http://www.statmt.org/wmt20/pdf/2020.wmt-1.101.pdf6""",
								    ]		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict		,	_UpperCAmelCase      :		List[str]      )  ->    Dict:





								'''simple docstring'''




								if self.config_name == "default":
												_lowerCAmelCase    :    Union[str, Any]										=      comet.load_from_checkpoint(comet.download_model("""wmt20-comet-da"""      )      )
								else:
												_lowerCAmelCase    :    Dict										=      comet.load_from_checkpoint(comet.download_model(self.config_name      )      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		Dict=None		,	_UpperCAmelCase      :		Tuple=False      )  ->    int:





								'''simple docstring'''




								if gpus is None:
												_lowerCAmelCase    :    Tuple										=      1 if torch.cuda.is_available() else 0
								_lowerCAmelCase    :    Optional[int]										=      {"""src""": sources, """mt""": predictions, """ref""": references}
								_lowerCAmelCase    :    List[Any]										=      [dict(zip(_UpperCAmelCase		,	_UpperCAmelCase      )      ) for t in zip(*data.values()      )]
								_lowerCAmelCase					,    _lowerCAmelCase    :    Optional[Any]										=      self.scorer.predict(_UpperCAmelCase		,	gpus=_UpperCAmelCase		,	progress_bar=_UpperCAmelCase      )
								return {"mean_score": mean_score, "scores": scores}

