




"""simple docstring"""
import comet  # From: unbabel-comet
import torch

import datasets


a_				:    Union[str, Any]										=     datasets.logging.get_logger(__name__)

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

a_				:    Optional[Any]										=     '''\
Crosslingual 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).
With 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.

See the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.
'''

a_				:    int										=     '''
COMET score.

Args:

`sources` (list of str): Source sentences
`predictions` (list of str): candidate translations
`references` (list of str): reference translations
`cuda` (bool): If set to True, runs COMET using GPU
`show_progress` (bool): Shows progress
`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.

Returns:
    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.
    `scores`: List of scores.

Examples:

    >>> comet_metric = datasets.load_metric(\'comet\')
    >>> # comet_metric = load_metric(\'comet\', \'wmt20-comet-da\')  # you can also choose which model to use
    >>> source = ["Dem Feuer konnte Einhalt geboten werden", "Schulen und Kindergärten wurden eröffnet."]
    >>> hypothesis = ["The fire could be stopped", "Schools and kindergartens were open"]
    >>> reference = ["They were able to control the fire.", "Schools and kindergartens opened"]
    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)
    >>> print([round(v, 2) for v in results["scores"]])
    [0.19, 0.92]
'''

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




					'''simple docstring'''


					def     snake_case_				(    self     ):
								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     snake_case_				(    self						,	__a     ):
								if self.config_name == "default":
											__lowerCamelCase					:	Any	       =				comet.load_from_checkpoint(comet.download_model('wmt20-comet-da'     )     )
								else:
											__lowerCamelCase					:	Union[str, Any]	       =				comet.load_from_checkpoint(comet.download_model(self.config_name     )     )





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









"""simple docstring"""
from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import TFLEDForConditionalGeneration, TFLEDModel

@require_tf
class __lowercase:




					'''simple docstring'''


					__a						:						Optional[Any]     		= LEDConfig
					__a						:						Dict     		= {}
					__a						:						int     		= 'gelu'
					def __init__(    self						,	__a						,	__a=13						,	__a=7						,	__a=True						,	__a=False						,	__a=99						,	__a=32						,	__a=2						,	__a=4						,	__a=37						,	__a=0.1						,	__a=0.1						,	__a=20						,	__a=2						,	__a=1						,	__a=0						,	__a=4						,	):
								__lowerCamelCase					:	str	       =				parent
								__lowerCamelCase					:	Optional[int]	       =				batch_size
								__lowerCamelCase					:	List[str]	       =				seq_length
								__lowerCamelCase					:	Optional[Any]	       =				is_training
								__lowerCamelCase					:	Dict	       =				use_labels
								__lowerCamelCase					:	Any	       =				vocab_size
								__lowerCamelCase					:	int	       =				hidden_size
								__lowerCamelCase					:	Any	       =				num_hidden_layers
								__lowerCamelCase					:	Union[str, Any]	       =				num_attention_heads
								__lowerCamelCase					:	List[str]	       =				intermediate_size

								__lowerCamelCase					:	str	       =				hidden_dropout_prob
								__lowerCamelCase					:	Union[str, Any]	       =				attention_probs_dropout_prob
								__lowerCamelCase					:	Any	       =				max_position_embeddings
								__lowerCamelCase					:	str	       =				eos_token_id
								__lowerCamelCase					:	str	       =				pad_token_id
								__lowerCamelCase					:	str	       =				bos_token_id
								__lowerCamelCase					:	int	       =				attention_window

								# `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
								# [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
								# returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
								# because its local attention only attends to `self.attention_window` and one before and one after
								__lowerCamelCase					:	Dict	       =				self.attention_window + 2

								# because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
								# the `test_attention_outputs` and `test_hidden_states_output` tests
								__lowerCamelCase					:	int	       =				(
								    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
								)
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	Tuple	       =				ids_tensor([self.batch_size, self.seq_length - 1]						,	self.vocab_size     )
								__lowerCamelCase					:	Optional[int]	       =				tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size     )						,	1     )
								__lowerCamelCase					:	Union[str, Any]	       =				tf.concat([input_ids, eos_tensor]						,	axis=1     )

								__lowerCamelCase					:	Any	       =				ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size     )

								__lowerCamelCase					:	Dict	       =				self.config_cls(
								    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_ids=[2]						,	bos_token_id=self.bos_token_id						,	pad_token_id=self.pad_token_id						,	decoder_start_token_id=self.pad_token_id						,	attention_window=self.attention_window						,	**self.config_updates						,	)
								__lowerCamelCase					:	int	       =				prepare_led_inputs_dict(__a						,	__a						,	__a     )
								__lowerCamelCase					:	Union[str, Any]	       =				tf.concat(
								    [tf.zeros_like(__a     )[:, :-1], tf.ones_like(__a     )[:, -1:]]						,	axis=-1						,	)
								__lowerCamelCase					:	Union[str, Any]	       =				global_attention_mask
								return config, inputs_dict





					def     snake_case_				(    self						,	__a						,	__a     ):
								__lowerCamelCase					:	Optional[int]	       =				TFLEDModel(config=__a     ).get_decoder()
								__lowerCamelCase					:	List[str]	       =				inputs_dict['input_ids']

								__lowerCamelCase					:	Dict	       =				input_ids[:1, :]
								__lowerCamelCase					:	Any	       =				inputs_dict['attention_mask'][:1, :]
								__lowerCamelCase					:	Tuple	       =				1

								# first forward pass
								__lowerCamelCase					:	List[str]	       =				model(__a						,	attention_mask=__a						,	use_cache=__a     )

								__lowerCamelCase       ,      __lowerCamelCase					:	int	       =				outputs.to_tuple()

								# create hypothetical next token and extent to next_input_ids
								__lowerCamelCase					:	Optional[int]	       =				ids_tensor((self.batch_size, 3)						,	config.vocab_size     )
								__lowerCamelCase					:	List[str]	       =				tf.cast(ids_tensor((self.batch_size, 3)						,	2     )						,	tf.inta     )

								# append to next input_ids and
								__lowerCamelCase					:	Any	       =				tf.concat([input_ids, next_tokens]						,	axis=-1     )
								__lowerCamelCase					:	Any	       =				tf.concat([attention_mask, next_attn_mask]						,	axis=-1     )

								__lowerCamelCase					:	Dict	       =				model(__a						,	attention_mask=__a     )[0]
								__lowerCamelCase					:	Optional[int]	       =				model(__a						,	attention_mask=__a						,	past_key_values=__a     )[0]

								self.parent.assertEqual(next_tokens.shape[1]						,	output_from_past.shape[1]     )

								# select random slice
								__lowerCamelCase					:	Any	       =				int(ids_tensor((1,)						,	output_from_past.shape[-1]     )     )
								__lowerCamelCase					:	int	       =				output_from_no_past[:, -3:, random_slice_idx]
								__lowerCamelCase					:	Union[str, Any]	       =				output_from_past[:, :, random_slice_idx]

								# test that outputs are equal for slice
								tf.debugging.assert_near(__a						,	__a						,	rtol=1E-3     )





def 						UpperCAmelCase				(			A__:    Any							,				A__:    Union[str, Any]							,				A__:    Dict							,				A__:    int=None							,				A__:    str=None							,				A__:    Tuple=None							,				A__:    List[Any]=None							,				)  ->     Optional[Any]:
			if attention_mask is None:
						__lowerCamelCase					:	Optional[int]	       =				tf.cast(tf.math.not_equal(A__							,				config.pad_token_id     )							,				tf.inta     )
			if decoder_attention_mask is None:
						__lowerCamelCase					:	Dict	       =				tf.concat(
						    [
						        tf.ones(decoder_input_ids[:, :1].shape							,				dtype=tf.inta     ),
						        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]							,				config.pad_token_id     )							,				tf.inta     ),
						    ]							,				axis=-1							,				)
			if head_mask is None:
						__lowerCamelCase					:	Dict	       =				tf.ones((config.encoder_layers, config.encoder_attention_heads)     )
			if decoder_head_mask is None:
						__lowerCamelCase					:	Tuple	       =				tf.ones((config.decoder_layers, config.decoder_attention_heads)     )
			return {
			    "input_ids": input_ids,
			    "attention_mask": attention_mask,
			    "decoder_input_ids": decoder_input_ids,
			    "decoder_attention_mask": decoder_attention_mask,
			    "head_mask": head_mask,
			    "decoder_head_mask": decoder_head_mask,
			}

@require_tf
class __lowercase(      lowercase__	,	lowercase__	,	unittest.TestCase				):




					'''simple docstring'''


					__a						:						Union[str, Any]     		= (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
					__a						:						Any     		= (TFLEDForConditionalGeneration,) if is_tf_available() else ()
					__a						:						str     		= (
					    {
					        'conversational': TFLEDForConditionalGeneration,
					        'feature-extraction': TFLEDModel,
					        'summarization': TFLEDForConditionalGeneration,
					        'text2text-generation': TFLEDForConditionalGeneration,
					        'translation': TFLEDForConditionalGeneration,
					    }
					    if is_tf_available()
					    else {}
					)
					__a						:						Union[str, Any]     		= True
					__a						:						str     		= False
					__a						:						Dict     		= False
					__a						:						List[Any]     		= False
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	int	       =				TFLEDModelTester(self     )
								__lowerCamelCase					:	List[str]	       =				ConfigTester(self						,	config_class=__a     )
					def     snake_case_				(    self     ):
								self.config_tester.run_common_tests()
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	Optional[int]	       =				self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.check_decoder_model_past_large_inputs(*__a     )
					def     snake_case_				(    self     ):
								__lowerCamelCase       ,      __lowerCamelCase					:	Optional[Any]	       =				self.model_tester.prepare_config_and_inputs_for_common()
								__lowerCamelCase					:	Optional[Any]	       =				tf.zeros_like(inputs_dict['attention_mask']     )
								__lowerCamelCase					:	Any	       =				2
								__lowerCamelCase					:	Dict	       =				tf.where(
								    tf.range(self.model_tester.seq_length     )[None, :] < num_global_attn_indices						,	1						,	inputs_dict['global_attention_mask']						,	)

								__lowerCamelCase					:	List[Any]	       =				True
								__lowerCamelCase					:	Tuple	       =				self.model_tester.seq_length
								__lowerCamelCase					:	Dict	       =				self.model_tester.encoder_seq_length

								def check_decoder_attentions_output(__a     ):
											__lowerCamelCase					:	List[Any]	       =				outputs.decoder_attentions
											self.assertEqual(len(__a     )						,	self.model_tester.num_hidden_layers     )
											self.assertListEqual(
											    list(decoder_attentions[0].shape[-3:]     )						,	[self.model_tester.num_attention_heads, seq_length, seq_length]						,	)

								def check_encoder_attentions_output(__a     ):
											__lowerCamelCase					:	List[str]	       =				[t.numpy() for t in outputs.encoder_attentions]
											__lowerCamelCase					:	Optional[Any]	       =				[t.numpy() for t in outputs.encoder_global_attentions]
											self.assertEqual(len(__a     )						,	self.model_tester.num_hidden_layers     )
											self.assertEqual(len(__a     )						,	self.model_tester.num_hidden_layers     )
											self.assertListEqual(
											    list(attentions[0].shape[-3:]     )						,	[self.model_tester.num_attention_heads, seq_length, seq_length]						,	)
											self.assertListEqual(
											    list(global_attentions[0].shape[-3:]     )						,	[self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]						,	)

								for model_class in self.all_model_classes:
											__lowerCamelCase					:	int	       =				True
											__lowerCamelCase					:	int	       =				False
											__lowerCamelCase					:	Tuple	       =				False
											__lowerCamelCase					:	List[Any]	       =				model_class(__a     )
											__lowerCamelCase					:	Optional[Any]	       =				model(self._prepare_for_class(__a						,	__a     )     )
											__lowerCamelCase					:	Tuple	       =				len(__a     )
											self.assertEqual(config.output_hidden_states						,	__a     )
											check_encoder_attentions_output(__a     )

											if self.is_encoder_decoder:
														__lowerCamelCase					:	Any	       =				model_class(__a     )
														__lowerCamelCase					:	List[str]	       =				model(self._prepare_for_class(__a						,	__a     )     )
														self.assertEqual(config.output_hidden_states						,	__a     )
														check_decoder_attentions_output(__a     )

											# Check that output attentions can also be changed via the config
											del inputs_dict["output_attentions"]
											__lowerCamelCase					:	Tuple	       =				True
											__lowerCamelCase					:	Dict	       =				model_class(__a     )
											__lowerCamelCase					:	List[Any]	       =				model(self._prepare_for_class(__a						,	__a     )     )
											self.assertEqual(config.output_hidden_states						,	__a     )
											check_encoder_attentions_output(__a     )

											# Check attention is always last and order is fine
											__lowerCamelCase					:	List[Any]	       =				True
											__lowerCamelCase					:	int	       =				True
											__lowerCamelCase					:	List[str]	       =				model_class(__a     )
											__lowerCamelCase					:	Dict	       =				model(self._prepare_for_class(__a						,	__a     )     )

											self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)						,	len(__a     )     )
											self.assertEqual(model.config.output_hidden_states						,	__a     )
											check_encoder_attentions_output(__a     )
					@unittest.skip('LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'     )
					def     snake_case_				(    self     ):
								pass





					def     snake_case_				(    self     ):
								# TODO: Head-masking not yet implement
								pass





def 						UpperCAmelCase				(			A__:    Union[str, Any]     )  ->     List[Any]:
			return tf.constant(A__							,				dtype=tf.intaa     )


a_				:    Tuple										=     1e-4






@slow
@require_tf
class __lowercase(      unittest.TestCase				):




					'''simple docstring'''


					def     snake_case_				(    self     ):
								__lowerCamelCase					:	Tuple	       =				TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'     ).led

								# change to intended input here
								__lowerCamelCase					:	Union[str, Any]	       =				_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]     )
								__lowerCamelCase					:	Dict	       =				_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]     )
								__lowerCamelCase					:	Optional[int]	       =				prepare_led_inputs_dict(model.config						,	__a						,	__a     )
								__lowerCamelCase					:	str	       =				model(**__a     )[0]
								__lowerCamelCase					:	Optional[int]	       =				(1, 1024, 768)
								self.assertEqual(output.shape						,	__a     )
								# change to expected output here
								__lowerCamelCase					:	int	       =				tf.convert_to_tensor(
								    [[2.3_050, 2.8_279, 0.6_531], [-1.8_457, -0.1_455, -3.5_661], [-1.0_186, 0.4_586, -2.2_043]]						,	)
								tf.debugging.assert_near(output[:, :3, :3]						,	__a						,	atol=1E-3     )





					def     snake_case_				(    self     ):
								__lowerCamelCase					:	List[str]	       =				TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'     )

								# change to intended input here
								__lowerCamelCase					:	Union[str, Any]	       =				_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]     )
								__lowerCamelCase					:	Dict	       =				_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]     )
								__lowerCamelCase					:	Optional[Any]	       =				prepare_led_inputs_dict(model.config						,	__a						,	__a     )
								__lowerCamelCase					:	Dict	       =				model(**__a     )[0]
								__lowerCamelCase					:	List[Any]	       =				(1, 1024, model.config.vocab_size)
								self.assertEqual(output.shape						,	__a     )
								# change to expected output here
								__lowerCamelCase					:	Optional[Any]	       =				tf.convert_to_tensor(
								    [[33.6_507, 6.4_572, 16.8_089], [5.8_739, -2.4_238, 11.2_902], [-3.2_139, -4.3_149, 4.2_783]]						,	)
								tf.debugging.assert_near(output[:, :3, :3]						,	__a						,	atol=1E-3						,	rtol=1E-3     )



