




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
	import os

	# The slow tests are often failing with OOM error on GPU
	# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
	# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
	_lowerCamelCase      :  List[str]      						=						"platform"

	import jax
	import jax.numpy as jnp

	from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
	    FlaxBlenderbotSmallForConditionalGeneration,
	    FlaxBlenderbotSmallModel,
	    shift_tokens_right,
	)




def 							_UpperCAmelCase  (UpperCamelCase_ : Dict      ,    UpperCamelCase_ : Optional[Any]      ,    UpperCamelCase_ : Optional[int]=None      ,    UpperCamelCase_ : List[str]=None      ,    UpperCamelCase_ : List[str]=None      ,    UpperCamelCase_ : Tuple=None      ,    UpperCamelCase_ : Tuple=None      ,    UpperCamelCase_ : int=None      ,    ):
				'''simple docstring'''
				if attention_mask is None:
								_lowerCAmelCase    :    List[str]										=      np.where(input_ids != config.pad_token_id      ,    1      ,    0       )
				if decoder_attention_mask is None:
								_lowerCAmelCase    :    Any										=      np.where(decoder_input_ids != config.pad_token_id      ,    1      ,    0       )
				if head_mask is None:
								_lowerCAmelCase    :    List[str]										=      np.ones((config.encoder_layers, config.encoder_attention_heads)       )
				if decoder_head_mask is None:
								_lowerCAmelCase    :    Dict										=      np.ones((config.decoder_layers, config.decoder_attention_heads)       )
				if cross_attn_head_mask is None:
								_lowerCAmelCase    :    Tuple										=      np.ones((config.decoder_layers, config.decoder_attention_heads)       )
				return {
				    "input_ids": input_ids,
				    "decoder_input_ids": decoder_input_ids,
				    "attention_mask": attention_mask,
				    "decoder_attention_mask": attention_mask,
				}
class        __snake_case  :



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





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      parent
								_lowerCAmelCase    :    Optional[Any]										=      batch_size
								_lowerCAmelCase    :    List[str]										=      seq_length
								_lowerCAmelCase    :    Tuple										=      is_training
								_lowerCAmelCase    :    List[str]										=      use_labels
								_lowerCAmelCase    :    Optional[int]										=      vocab_size
								_lowerCAmelCase    :    int										=      hidden_size
								_lowerCAmelCase    :    Union[str, Any]										=      num_hidden_layers
								_lowerCAmelCase    :    List[str]										=      num_attention_heads
								_lowerCAmelCase    :    Any										=      intermediate_size
								_lowerCAmelCase    :    Tuple										=      hidden_act
								_lowerCAmelCase    :    Dict										=      hidden_dropout_prob
								_lowerCAmelCase    :    Optional[Any]										=      attention_probs_dropout_prob
								_lowerCAmelCase    :    Optional[Any]										=      max_position_embeddings
								_lowerCAmelCase    :    List[Any]										=      eos_token_id
								_lowerCAmelCase    :    Optional[Any]										=      pad_token_id
								_lowerCAmelCase    :    str										=      bos_token_id
								_lowerCAmelCase    :    Dict										=      initializer_range



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





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      np.clip(ids_tensor([self.batch_size, self.seq_length - 1]		,	self.vocab_size      )		,	3		,	self.vocab_size      )
								_lowerCAmelCase    :    Any										=      np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)		,	dtype=np.intaa      ))		,	-1      )

								_lowerCAmelCase    :    int										=      shift_tokens_right(_UpperCAmelCase		,	1		,	2      )

								_lowerCAmelCase    :    int										=      BlenderbotSmallConfig(
								    vocab_size=self.vocab_size		,	d_model=self.hidden_size		,	encoder_layers=self.num_hidden_layers		,	decoder_layers=self.num_hidden_layers		,	encoder_attention_heads=self.num_attention_heads		,	decoder_attention_heads=self.num_attention_heads		,	encoder_ffn_dim=self.intermediate_size		,	decoder_ffn_dim=self.intermediate_size		,	dropout=self.hidden_dropout_prob		,	attention_dropout=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	eos_token_id=self.eos_token_id		,	bos_token_id=self.bos_token_id		,	pad_token_id=self.pad_token_id		,	initializer_range=self.initializer_range		,	use_cache=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    List[Any]										=      prepare_blenderbot_inputs_dict(_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase      )
								return config, inputs_dict



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





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    Dict										=      self.prepare_config_and_inputs()
								return config, inputs_dict



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





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      20
								_lowerCAmelCase    :    Optional[int]										=      model_class_name(_UpperCAmelCase      )

								_lowerCAmelCase    :    Union[str, Any]										=      model.encode(inputs_dict["""input_ids"""]      )

								_lowerCAmelCase					,    _lowerCAmelCase    :    List[Any]										=      (
								    inputs_dict["""decoder_input_ids"""],
								    inputs_dict["""decoder_attention_mask"""],
								)

								_lowerCAmelCase    :    str										=      model.init_cache(decoder_input_ids.shape[0]		,	_UpperCAmelCase		,	_UpperCAmelCase      )
								_lowerCAmelCase    :    Any										=      jnp.ones((decoder_input_ids.shape[0], max_decoder_length)		,	dtype="""i4"""      )

								_lowerCAmelCase    :    Union[str, Any]										=      jnp.broadcast_to(
								    jnp.arange(decoder_input_ids.shape[-1] - 1      )[None, :]		,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,	)
								_lowerCAmelCase    :    List[str]										=      model.decode(
								    decoder_input_ids[:, :-1]		,	_UpperCAmelCase		,	decoder_attention_mask=_UpperCAmelCase		,	past_key_values=_UpperCAmelCase		,	decoder_position_ids=_UpperCAmelCase		,	)

								_lowerCAmelCase    :    Union[str, Any]										=      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,	dtype="""i4"""      )
								_lowerCAmelCase    :    Union[str, Any]										=      model.decode(
								    decoder_input_ids[:, -1:]		,	_UpperCAmelCase		,	decoder_attention_mask=_UpperCAmelCase		,	past_key_values=outputs_cache.past_key_values		,	decoder_position_ids=_UpperCAmelCase		,	)

								_lowerCAmelCase    :    str										=      model.decode(_UpperCAmelCase		,	_UpperCAmelCase      )

								_lowerCAmelCase    :    Optional[int]										=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
								self.parent.assertTrue(diff < 1E-3		,	msg=f"Max diff is {diff}"      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      20
								_lowerCAmelCase    :    str										=      model_class_name(_UpperCAmelCase      )

								_lowerCAmelCase    :    str										=      model.encode(inputs_dict["""input_ids"""]      )

								_lowerCAmelCase					,    _lowerCAmelCase    :    List[str]										=      (
								    inputs_dict["""decoder_input_ids"""],
								    inputs_dict["""decoder_attention_mask"""],
								)

								_lowerCAmelCase    :    Optional[Any]										=      jnp.concatenate(
								    [
								        decoder_attention_mask,
								        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])      ),
								    ]		,	axis=-1		,	)

								_lowerCAmelCase    :    List[Any]										=      model.init_cache(decoder_input_ids.shape[0]		,	_UpperCAmelCase		,	_UpperCAmelCase      )
								_lowerCAmelCase    :    Tuple										=      jnp.broadcast_to(
								    jnp.arange(decoder_input_ids.shape[-1] - 1      )[None, :]		,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,	)

								_lowerCAmelCase    :    str										=      model.decode(
								    decoder_input_ids[:, :-1]		,	_UpperCAmelCase		,	decoder_attention_mask=_UpperCAmelCase		,	past_key_values=_UpperCAmelCase		,	decoder_position_ids=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    Tuple										=      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,	dtype="""i4"""      )
								_lowerCAmelCase    :    Optional[int]										=      model.decode(
								    decoder_input_ids[:, -1:]		,	_UpperCAmelCase		,	past_key_values=outputs_cache.past_key_values		,	decoder_attention_mask=_UpperCAmelCase		,	decoder_position_ids=_UpperCAmelCase		,	)

								_lowerCAmelCase    :    str										=      model.decode(_UpperCAmelCase		,	_UpperCAmelCase		,	decoder_attention_mask=_UpperCAmelCase      )

								_lowerCAmelCase    :    str										=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
								self.parent.assertTrue(diff < 1E-3		,	msg=f"Max diff is {diff}"      )






@require_flax
class        __snake_case  (unittest.TestCase		):
				lowerCAmelCase__				=     9_9



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





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      np.array(
								    [
								        [71, 82, 18, 33, 46, 91, 2],
								        [68, 34, 26, 58, 30, 82, 2],
								        [5, 97, 17, 39, 94, 40, 2],
								        [76, 83, 94, 25, 70, 78, 2],
								        [87, 59, 41, 35, 48, 66, 2],
								        [55, 13, 16, 58, 5, 2, 1],  # note padding
								        [64, 27, 31, 51, 12, 75, 2],
								        [52, 64, 86, 17, 83, 39, 2],
								        [48, 61, 9, 24, 71, 82, 2],
								        [26, 1, 60, 48, 22, 13, 2],
								        [21, 5, 62, 28, 14, 76, 2],
								        [45, 98, 37, 86, 59, 48, 2],
								        [70, 70, 50, 9, 28, 0, 2],
								    ]		,	dtype=np.intaa		,	)

								_lowerCAmelCase    :    Union[str, Any]										=      input_ids.shape[0]
								_lowerCAmelCase    :    str										=      BlenderbotSmallConfig(
								    vocab_size=self.vocab_size		,	d_model=24		,	encoder_layers=2		,	decoder_layers=2		,	encoder_attention_heads=2		,	decoder_attention_heads=2		,	encoder_ffn_dim=32		,	decoder_ffn_dim=32		,	max_position_embeddings=48		,	eos_token_id=2		,	pad_token_id=1		,	bos_token_id=0		,	)
								return config, input_ids, batch_size



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase					,    _lowerCAmelCase    :    Union[str, Any]										=      self._get_config_and_data()
								_lowerCAmelCase    :    Dict										=      FlaxBlenderbotSmallForConditionalGeneration(_UpperCAmelCase      )
								_lowerCAmelCase    :    Tuple										=      lm_model(input_ids=_UpperCAmelCase      )
								_lowerCAmelCase    :    Dict										=      (batch_size, input_ids.shape[1], config.vocab_size)
								self.assertEqual(outputs["""logits"""].shape		,	_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      BlenderbotSmallConfig(
								    vocab_size=self.vocab_size		,	d_model=14		,	encoder_layers=2		,	decoder_layers=2		,	encoder_attention_heads=2		,	decoder_attention_heads=2		,	encoder_ffn_dim=8		,	decoder_ffn_dim=8		,	max_position_embeddings=48		,	)
								_lowerCAmelCase    :    Any										=      FlaxBlenderbotSmallForConditionalGeneration(_UpperCAmelCase      )
								_lowerCAmelCase    :    Tuple										=      np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]		,	dtype=np.intaa      )
								_lowerCAmelCase    :    List[Any]										=      np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]		,	dtype=np.intaa      )
								_lowerCAmelCase    :    Optional[Any]										=      lm_model(input_ids=_UpperCAmelCase		,	decoder_input_ids=_UpperCAmelCase      )
								_lowerCAmelCase    :    Any										=      (*summary.shape, config.vocab_size)
								self.assertEqual(outputs["""logits"""].shape		,	_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]		,	dtype=np.intaa      )
								_lowerCAmelCase    :    List[Any]										=      shift_tokens_right(_UpperCAmelCase		,	1		,	2      )
								_lowerCAmelCase    :    int										=      np.equal(_UpperCAmelCase		,	1      ).astype(np.floataa      ).sum()
								_lowerCAmelCase    :    List[str]										=      np.equal(_UpperCAmelCase		,	1      ).astype(np.floataa      ).sum()
								self.assertEqual(shifted.shape		,	input_ids.shape      )
								self.assertEqual(_UpperCAmelCase		,	n_pad_before - 1      )
								self.assertTrue(np.equal(shifted[:, 0]		,	2      ).all()      )






@require_flax
class        __snake_case  (_a						,	unittest.TestCase						,	_a		):
				lowerCAmelCase__				=     True
				lowerCAmelCase__				=     (
				    (
				        FlaxBlenderbotSmallModel,
				        FlaxBlenderbotSmallForConditionalGeneration,
				    )
				    if is_flax_available()
				    else ()
				)
				lowerCAmelCase__				=     (FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()



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





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      FlaxBlenderbotSmallModelTester(self      )



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





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    int										=      self.model_tester.prepare_config_and_inputs()
								for model_class in self.all_model_classes:
												self.model_tester.check_use_cache_forward(_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    Dict										=      self.model_tester.prepare_config_and_inputs()
								for model_class in self.all_model_classes:
												self.model_tester.check_use_cache_forward_with_attn_mask(_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    List[str]										=      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__      ):
																_lowerCAmelCase    :    Tuple										=      self._prepare_for_class(_UpperCAmelCase		,	_UpperCAmelCase      )
																_lowerCAmelCase    :    Optional[Any]										=      model_class(_UpperCAmelCase      )

																@jax.jit
																def encode_jitted(_UpperCAmelCase      :		Optional[Any]		,	_UpperCAmelCase      :		Dict=None		,	**_UpperCAmelCase      :		int      ):
																				return model.encode(input_ids=_UpperCAmelCase		,	attention_mask=_UpperCAmelCase      )

																with self.subTest("""JIT Enabled"""      ):
																				_lowerCAmelCase    :    Optional[int]										=      encode_jitted(**_UpperCAmelCase      ).to_tuple()

																with self.subTest("""JIT Disabled"""      ):
																				with jax.disable_jit():
																								_lowerCAmelCase    :    Optional[int]										=      encode_jitted(**_UpperCAmelCase      ).to_tuple()

																self.assertEqual(len(_UpperCAmelCase      )		,	len(_UpperCAmelCase      )      )
																for jitted_output, output in zip(_UpperCAmelCase		,	_UpperCAmelCase      ):
																				self.assertEqual(jitted_output.shape		,	output.shape      )



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





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    List[Any]										=      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__      ):
																_lowerCAmelCase    :    Union[str, Any]										=      model_class(_UpperCAmelCase      )
																_lowerCAmelCase    :    int										=      model.encode(inputs_dict["""input_ids"""]		,	inputs_dict["""attention_mask"""]      )

																_lowerCAmelCase    :    Any										=      {
																    """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
																    """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
																    """encoder_outputs""": encoder_outputs,
																}

																@jax.jit
																def decode_jitted(_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		Optional[int]		,	_UpperCAmelCase      :		Union[str, Any]      ):
																				return model.decode(
																				    decoder_input_ids=_UpperCAmelCase		,	decoder_attention_mask=_UpperCAmelCase		,	encoder_outputs=_UpperCAmelCase		,	)

																with self.subTest("""JIT Enabled"""      ):
																				_lowerCAmelCase    :    Dict										=      decode_jitted(**_UpperCAmelCase      ).to_tuple()

																with self.subTest("""JIT Disabled"""      ):
																				with jax.disable_jit():
																								_lowerCAmelCase    :    List[str]										=      decode_jitted(**_UpperCAmelCase      ).to_tuple()

																self.assertEqual(len(_UpperCAmelCase      )		,	len(_UpperCAmelCase      )      )
																for jitted_output, output in zip(_UpperCAmelCase		,	_UpperCAmelCase      ):
																				self.assertEqual(jitted_output.shape		,	output.shape      )



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    str:





								'''simple docstring'''




								for model_class_name in self.all_model_classes:
												_lowerCAmelCase    :    str										=      model_class_name.from_pretrained("""facebook/blenderbot_small-90M"""      )
												# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
												_lowerCAmelCase    :    Union[str, Any]										=      np.ones((1, 1)      ) * model.config.eos_token_id
												_lowerCAmelCase    :    Any										=      model(_UpperCAmelCase      )
												self.assertIsNotNone(_UpperCAmelCase      )







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}

