




import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
			from huggingface_hub import snapshot_download
			from pyctcdecode import BeamSearchDecoderCTC

			from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
			from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
			from transformers import WavaVecaForCTC






@require_pyctcdecode
class lowerCAmelCase__			(  unittest.TestCase		):






			'''simple docstring'''



			def       A_					(						self						):
							_lowerCamelCase  :    Union[str, Any]					=					"""| <pad> <unk> <s> </s> a b c d e f g h i j k""".split()
							_lowerCamelCase  :    int					=					dict(zip(lowercase						,   range(len(lowercase						)						)						)						)

							_lowerCamelCase  :    List[Any]					=					{
							    """unk_token""": """<unk>""",
							    """bos_token""": """<s>""",
							    """eos_token""": """</s>""",
							}
							_lowerCamelCase  :    Tuple					=					{
							    """feature_size""": 1,
							    """padding_value""": 0.0,
							    """sampling_rate""": 16000,
							    """return_attention_mask""": False,
							    """do_normalize""": True,
							}

							_lowerCamelCase  :    List[Any]					=					tempfile.mkdtemp()
							_lowerCamelCase  :    str					=					os.path.join(self.tmpdirname						,   VOCAB_FILES_NAMES['vocab_file']						)
							_lowerCamelCase  :    Tuple					=					os.path.join(self.tmpdirname						,   lowercase						)
							with open(self.vocab_file						,   'w'						,   encoding='utf-8'						) as fp:
											fp.write(json.dumps(lowercase						) + '\n'						)

							with open(self.feature_extraction_file						,   'w'						,   encoding='utf-8'						) as fp:
											fp.write(json.dumps(lowercase						) + '\n'						)

							# load decoder from hub
							_lowerCamelCase  :    List[str]					=					"""hf-internal-testing/ngram-beam-search-decoder"""
			def       A_					(						self						,   **lowercase						):
							_lowerCamelCase  :    Union[str, Any]					=					self.add_kwargs_tokens_map.copy()
							kwargs.update(lowercase						)
							return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname						,   **lowercase						)
			def       A_					(						self						,   **lowercase						):
							return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname						,   **lowercase						)
			def       A_					(						self						,   **lowercase						):
							return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name						,   **lowercase						)
			def       A_					(						self						):
							shutil.rmtree(self.tmpdirname						)
			def       A_					(						self						):
							_lowerCamelCase  :    Dict					=					self.get_tokenizer()
							_lowerCamelCase  :    Optional[int]					=					self.get_feature_extractor()
							_lowerCamelCase  :    int					=					self.get_decoder()

							_lowerCamelCase  :    List[str]					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							processor.save_pretrained(self.tmpdirname						)
							_lowerCamelCase  :    Dict					=					WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname						)

							# tokenizer
							self.assertEqual(processor.tokenizer.get_vocab()						,   tokenizer.get_vocab()						)
							self.assertIsInstance(processor.tokenizer						,   lowercase						)

							# feature extractor
							self.assertEqual(processor.feature_extractor.to_json_string()						,   feature_extractor.to_json_string()						)
							self.assertIsInstance(processor.feature_extractor						,   lowercase						)

							# decoder
							self.assertEqual(processor.decoder._alphabet.labels						,   decoder._alphabet.labels						)
							self.assertEqual(
							    processor.decoder.model_container[decoder._model_key]._unigram_set						,   decoder.model_container[decoder._model_key]._unigram_set						,   )
							self.assertIsInstance(processor.decoder						,   lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    str					=					WavaVecaProcessorWithLM(
							    tokenizer=self.get_tokenizer()						,   feature_extractor=self.get_feature_extractor()						,   decoder=self.get_decoder()						)
							processor.save_pretrained(self.tmpdirname						)

							# make sure that error is thrown when decoder alphabet doesn't match
							_lowerCamelCase  :    Dict					=					WavaVecaProcessorWithLM.from_pretrained(
							    self.tmpdirname						,   alpha=5.0						,   beta=3.0						,   score_boundary=-7.0						,   unk_score_offset=3						)

							# decoder
							self.assertEqual(processor.language_model.alpha						,   5.0						)
							self.assertEqual(processor.language_model.beta						,   3.0						)
							self.assertEqual(processor.language_model.score_boundary						,   -7.0						)
							self.assertEqual(processor.language_model.unk_score_offset						,   3						)
			def       A_					(						self						):
							_lowerCamelCase  :    int					=					self.get_tokenizer()
							# add token to trigger raise
							tokenizer.add_tokens(['xx']						)
							with self.assertRaisesRegex(lowercase						,   'include'						):
											WavaVecaProcessorWithLM(
											    tokenizer=lowercase						,   feature_extractor=self.get_feature_extractor()						,   decoder=self.get_decoder()						)
			def       A_					(						self						):
							_lowerCamelCase  :    Optional[Any]					=					self.get_feature_extractor()
							_lowerCamelCase  :    Union[str, Any]					=					self.get_tokenizer()
							_lowerCamelCase  :    Any					=					self.get_decoder()

							_lowerCamelCase  :    Any					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							_lowerCamelCase  :    Tuple					=					floats_list((3, 1000)						)

							_lowerCamelCase  :    Any					=					feature_extractor(lowercase						,   return_tensors='np'						)
							_lowerCamelCase  :    int					=					processor(lowercase						,   return_tensors='np'						)

							for key in input_feat_extract.keys():
											self.assertAlmostEqual(input_feat_extract[key].sum()						,   input_processor[key].sum()						,   delta=1E-2						)
			def       A_					(						self						):
							_lowerCamelCase  :    Optional[int]					=					self.get_feature_extractor()
							_lowerCamelCase  :    Union[str, Any]					=					self.get_tokenizer()
							_lowerCamelCase  :    int					=					self.get_decoder()

							_lowerCamelCase  :    List[str]					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							_lowerCamelCase  :    Optional[int]					=					"""This is a test string"""

							_lowerCamelCase  :    List[str]					=					processor(text=lowercase						)

							_lowerCamelCase  :    Any					=					tokenizer(lowercase						)

							for key in encoded_tok.keys():
											self.assertListEqual(encoded_tok[key]						,   encoded_processor[key]						)
			def       A_					(						self						,   lowercase=(2, 10, 16)						,   lowercase=77						):
							np.random.seed(lowercase						)
							return np.random.rand(*lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    str					=					self.get_feature_extractor()
							_lowerCamelCase  :    Any					=					self.get_tokenizer()
							_lowerCamelCase  :    Optional[Any]					=					self.get_decoder()

							_lowerCamelCase  :    Optional[int]					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							_lowerCamelCase  :    Tuple					=					self._get_dummy_logits(shape=(10, 16)						,   seed=13						)

							_lowerCamelCase  :    Optional[Any]					=					processor.decode(lowercase						)

							_lowerCamelCase  :    List[Any]					=					decoder.decode_beams(lowercase						)[0]

							self.assertEqual(decoded_decoder[0]						,   decoded_processor.text						)
							self.assertEqual('</s> <s> </s>'						,   decoded_processor.text						)
							self.assertEqual(decoded_decoder[-2]						,   decoded_processor.logit_score						)
							self.assertEqual(decoded_decoder[-1]						,   decoded_processor.lm_score						)
			@parameterized.expand([[None], ['fork'], ['spawn']]						)
			def       A_					(						self						,   lowercase						):
							_lowerCamelCase  :    List[str]					=					self.get_feature_extractor()
							_lowerCamelCase  :    List[str]					=					self.get_tokenizer()
							_lowerCamelCase  :    str					=					self.get_decoder()

							_lowerCamelCase  :    str					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							_lowerCamelCase  :    str					=					self._get_dummy_logits()

							# note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
							#       otherwise, the LM won't be available to the pool's sub-processes.
							# manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
							if pool_context is None:
											_lowerCamelCase  :    List[str]					=					processor.batch_decode(lowercase						)
							else:
											with get_context(lowercase						).Pool() as pool:
															_lowerCamelCase  :    Tuple					=					processor.batch_decode(lowercase						,   lowercase						)

							_lowerCamelCase  :    str					=					list(lowercase						)

							with get_context('fork'						).Pool() as p:
											_lowerCamelCase  :    List[str]					=					decoder.decode_beams_batch(lowercase						,   lowercase						)

							_lowerCamelCase  :    Dict					=					[], [], []
							for beams in decoded_beams:
											texts_decoder.append(beams[0][0]						)
											logit_scores_decoder.append(beams[0][-2]						)
											lm_scores_decoder.append(beams[0][-1]						)

							self.assertListEqual(lowercase						,   decoded_processor.text						)
							self.assertListEqual(['<s> <s> </s>', '<s> <s> <s>']						,   decoded_processor.text						)
							self.assertListEqual(lowercase						,   decoded_processor.logit_score						)
							self.assertListEqual(lowercase						,   decoded_processor.lm_score						)
			def       A_					(						self						):
							_lowerCamelCase  :    Any					=					self.get_feature_extractor()
							_lowerCamelCase  :    Optional[int]					=					self.get_tokenizer()
							_lowerCamelCase  :    Any					=					self.get_decoder()

							_lowerCamelCase  :    List[Any]					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							_lowerCamelCase  :    int					=					self._get_dummy_logits()

							_lowerCamelCase  :    Optional[Any]					=					15
							_lowerCamelCase  :    Union[str, Any]					=					-20.0
							_lowerCamelCase  :    List[Any]					=					-4.0

							_lowerCamelCase  :    List[str]					=					processor.batch_decode(
							    lowercase						,   beam_width=lowercase						,   beam_prune_logp=lowercase						,   token_min_logp=lowercase						,   )
							_lowerCamelCase  :    str					=					decoded_processor_out.text

							_lowerCamelCase  :    int					=					list(lowercase						)

							with get_context('fork'						).Pool() as pool:
											_lowerCamelCase  :    List[Any]					=					decoder.decode_beams_batch(
											    lowercase						,   lowercase						,   beam_width=lowercase						,   beam_prune_logp=lowercase						,   token_min_logp=lowercase						,   )

							_lowerCamelCase  :    str					=					[d[0][0] for d in decoded_decoder_out]
							_lowerCamelCase  :    List[Any]					=					[d[0][2] for d in decoded_decoder_out]
							_lowerCamelCase  :    Optional[Any]					=					[d[0][3] for d in decoded_decoder_out]

							self.assertListEqual(lowercase						,   lowercase						)
							self.assertListEqual(['</s> <s> <s>', '<s> <s> <s>']						,   lowercase						)

							self.assertTrue(np.array_equal(lowercase						,   decoded_processor_out.logit_score						)						)
							self.assertTrue(np.allclose([-20.0_54, -18.4_47]						,   lowercase						,   atol=1E-3						)						)

							self.assertTrue(np.array_equal(lowercase						,   decoded_processor_out.lm_score						)						)
							self.assertTrue(np.allclose([-15.5_54, -13.94_74]						,   lowercase						,   atol=1E-3						)						)
			def       A_					(						self						):
							_lowerCamelCase  :    Union[str, Any]					=					self.get_feature_extractor()
							_lowerCamelCase  :    Dict					=					self.get_tokenizer()
							_lowerCamelCase  :    List[str]					=					self.get_decoder()

							_lowerCamelCase  :    Union[str, Any]					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							_lowerCamelCase  :    Any					=					self._get_dummy_logits()

							_lowerCamelCase  :    str					=					2.0
							_lowerCamelCase  :    Tuple					=					5.0
							_lowerCamelCase  :    Dict					=					-20.0
							_lowerCamelCase  :    List[Any]					=					True

							_lowerCamelCase  :    Union[str, Any]					=					processor.batch_decode(
							    lowercase						,   alpha=lowercase						,   beta=lowercase						,   unk_score_offset=lowercase						,   lm_score_boundary=lowercase						,   )
							_lowerCamelCase  :    List[str]					=					decoded_processor_out.text

							_lowerCamelCase  :    Optional[Any]					=					list(lowercase						)
							decoder.reset_params(
							    alpha=lowercase						,   beta=lowercase						,   unk_score_offset=lowercase						,   lm_score_boundary=lowercase						,   )

							with get_context('fork'						).Pool() as pool:
											_lowerCamelCase  :    int					=					decoder.decode_beams_batch(
											    lowercase						,   lowercase						,   )

							_lowerCamelCase  :    Optional[Any]					=					[d[0][0] for d in decoded_decoder_out]

							self.assertListEqual(lowercase						,   lowercase						)
							self.assertListEqual(['<s> </s> <s> </s> </s>', '</s> </s> <s> </s> </s>']						,   lowercase						)
							_lowerCamelCase  :    Optional[Any]					=					processor.decoder.model_container[processor.decoder._model_key]
							self.assertEqual(lm_model.alpha						,   2.0						)
							self.assertEqual(lm_model.beta						,   5.0						)
							self.assertEqual(lm_model.unk_score_offset						,   -20.0						)
							self.assertEqual(lm_model.score_boundary						,   lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    Union[str, Any]					=					WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'						)

							_lowerCamelCase  :    int					=					processor.decoder.model_container[processor.decoder._model_key]
							_lowerCamelCase  :    int					=					Path(language_model._kenlm_model.path.decode('utf-8'						)						).parent.parent.absolute()

							_lowerCamelCase  :    Any					=					os.listdir(lowercase						)
							_lowerCamelCase  :    List[Any]					=					["""alphabet.json""", """language_model"""]

							downloaded_decoder_files.sort()
							expected_decoder_files.sort()

							# test that only decoder relevant files from
							# https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
							# are downloaded and none of the rest (e.g. README.md, ...)
							self.assertListEqual(lowercase						,   lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    Tuple					=					snapshot_download('hf-internal-testing/processor_with_lm'						)

							_lowerCamelCase  :    Tuple					=					WavaVecaProcessorWithLM.from_pretrained(lowercase						)

							_lowerCamelCase  :    Optional[int]					=					processor.decoder.model_container[processor.decoder._model_key]
							_lowerCamelCase  :    Optional[Any]					=					Path(language_model._kenlm_model.path.decode('utf-8'						)						).parent.parent.absolute()

							_lowerCamelCase  :    Any					=					os.listdir(lowercase						)
							_lowerCamelCase  :    List[str]					=					os.listdir(lowercase						)

							local_decoder_files.sort()
							expected_decoder_files.sort()

							# test that both decoder form hub and local files in cache are the same
							self.assertListEqual(lowercase						,   lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    List[str]					=					WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'						)
							_lowerCamelCase  :    Any					=					AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm'						)

							_lowerCamelCase  :    Union[str, Any]					=					floats_list((3, 1000)						)

							_lowerCamelCase  :    Tuple					=					processor_wavaveca(lowercase						,   return_tensors='np'						)
							_lowerCamelCase  :    Tuple					=					processor_auto(lowercase						,   return_tensors='np'						)

							for key in input_wavaveca.keys():
											self.assertAlmostEqual(input_wavaveca[key].sum()						,   input_auto[key].sum()						,   delta=1E-2						)

							_lowerCamelCase  :    str					=					self._get_dummy_logits()

							_lowerCamelCase  :    int					=					processor_wavaveca.batch_decode(lowercase						)
							_lowerCamelCase  :    Tuple					=					processor_auto.batch_decode(lowercase						)

							self.assertListEqual(decoded_wavaveca.text						,   decoded_auto.text						)
			def       A_					(						self						):
							_lowerCamelCase  :    List[Any]					=					self.get_feature_extractor()
							_lowerCamelCase  :    Dict					=					self.get_tokenizer()
							_lowerCamelCase  :    Optional[Any]					=					self.get_decoder()

							_lowerCamelCase  :    Dict					=					WavaVecaProcessorWithLM(tokenizer=lowercase						,   feature_extractor=lowercase						,   decoder=lowercase						)

							self.assertListEqual(
							    processor.model_input_names						,   feature_extractor.model_input_names						,   msg='`processor` and `feature_extractor` model input names do not match'						,   )
			@staticmethod
			def       A_					(						lowercase						,   lowercase						):
							_lowerCamelCase  :    List[Any]					=					[d[key] for d in offsets]
							return retrieved_list
			def       A_					(						self						):
							_lowerCamelCase  :    int					=					WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'						)
							_lowerCamelCase  :    int					=					self._get_dummy_logits()[0]

							_lowerCamelCase  :    List[Any]					=					processor.decode(lowercase						,   output_word_offsets=lowercase						)
							# check Wav2Vec2CTCTokenizerOutput keys for word
							self.assertEqual(len(outputs.keys()						)						,   4						)
							self.assertTrue('text' in outputs						)
							self.assertTrue('word_offsets' in outputs						)
							self.assertTrue(isinstance(lowercase						,   lowercase						)						)

							self.assertEqual(' '.join(self.get_from_offsets(outputs['word_offsets']						,   'word'						)						)						,   outputs.text						)
							self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,   'word'						)						,   ['<s>', '<s>', '</s>']						)
							self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,   'start_offset'						)						,   [0, 2, 4]						)
							self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,   'end_offset'						)						,   [1, 3, 5]						)
			def       A_					(						self						):
							_lowerCamelCase  :    Dict					=					WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'						)
							_lowerCamelCase  :    Tuple					=					self._get_dummy_logits()

							_lowerCamelCase  :    Optional[int]					=					processor.batch_decode(lowercase						,   output_word_offsets=lowercase						)

							# check Wav2Vec2CTCTokenizerOutput keys for word
							self.assertEqual(len(outputs.keys()						)						,   4						)
							self.assertTrue('text' in outputs						)
							self.assertTrue('word_offsets' in outputs						)
							self.assertTrue(isinstance(lowercase						,   lowercase						)						)

							self.assertListEqual(
							    [' '.join(self.get_from_offsets(lowercase						,   'word'						)						) for o in outputs['word_offsets']]						,   outputs.text						)
							self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,   'word'						)						,   ['<s>', '<s>', '</s>']						)
							self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,   'start_offset'						)						,   [0, 2, 4]						)
							self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,   'end_offset'						)						,   [1, 3, 5]						)






			@slow
			@require_torch
			@require_torchaudio
			def       A_					(						self						):
							import torch

							_lowerCamelCase  :    Tuple					=					load_dataset('common_voice'						,   'en'						,   split='train'						,   streaming=lowercase						)
							_lowerCamelCase  :    str					=					ds.cast_column('audio'						,   datasets.Audio(sampling_rate=16000						)						)
							_lowerCamelCase  :    Optional[int]					=					iter(lowercase						)
							_lowerCamelCase  :    Tuple					=					next(lowercase						)

							_lowerCamelCase  :    Any					=					AutoProcessor.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'						)
							_lowerCamelCase  :    str					=					WavaVecaForCTC.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'						)

							# compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
							_lowerCamelCase  :    Tuple					=					processor(sample['audio']['array']						,   return_tensors='pt'						).input_values

							with torch.no_grad():
											_lowerCamelCase  :    str					=					model(lowercase						).logits.cpu().numpy()

							_lowerCamelCase  :    List[str]					=					processor.decode(logits[0]						,   output_word_offsets=lowercase						)

							_lowerCamelCase  :    str					=					model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
							_lowerCamelCase  :    Union[str, Any]					=					[
							    {
							        """start_time""": d["""start_offset"""] * time_offset,
							        """end_time""": d["""end_offset"""] * time_offset,
							        """word""": d["""word"""],
							    }
							    for d in output["""word_offsets"""]
							]

							_lowerCamelCase  :    str					=					"""WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL"""

							# output words
							self.assertEqual(' '.join(self.get_from_offsets(lowercase						,   'word'						)						)						,   lowercase						)
							self.assertEqual(' '.join(self.get_from_offsets(lowercase						,   'word'						)						)						,   output.text						)

							# output times
							_lowerCamelCase  :    Any					=					torch.tensor(self.get_from_offsets(lowercase						,   'start_time'						)						)
							_lowerCamelCase  :    Dict					=					torch.tensor(self.get_from_offsets(lowercase						,   'end_time'						)						)

							# fmt: off
							_lowerCamelCase  :    Any					=					torch.tensor([1.41_99, 1.65_99, 2.25_99, 3.0, 3.24, 3.59_99, 3.79_99, 4.09_99, 4.26, 4.94, 5.28, 5.65_99, 5.78, 5.94, 6.32, 6.53_99, 6.65_99]						)
							_lowerCamelCase  :    str					=					torch.tensor([1.53_99, 1.89_99, 2.9, 3.16, 3.53_99, 3.72, 4.01_99, 4.17_99, 4.76, 5.15_99, 5.55_99, 5.69_99, 5.86, 6.19_99, 6.38, 6.61_99, 6.94]						)
							# fmt: on

							self.assertTrue(torch.allclose(lowercase						,   lowercase						,   atol=0.01						)						)
							self.assertTrue(torch.allclose(lowercase						,   lowercase						,   atol=0.01						)						)





"""simple docstring"""






import comet  # From: unbabel-comet
import torch

import datasets


lowercase__  		=		datasets.logging.get_logger(__name__)

lowercase__  		=		"""\
@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\",
}
"""

lowercase__  		=		"""\
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.
"""

lowercase__  		=		"""
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 lowerCAmelCase__			(  datasets.Metric		):






			'''simple docstring'''



			def       A_					(						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       A_					(						self						,   lowercase						):
							if self.config_name == "default":
											_lowerCamelCase  :    List[Any]					=					comet.load_from_checkpoint(comet.download_model('wmt20-comet-da'						)						)
							else:
											_lowerCamelCase  :    Tuple					=					comet.load_from_checkpoint(comet.download_model(self.config_name						)						)






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