


import comet  # From: unbabel-comet
import torch

import datasets


A   :					List[str]	=      datasets.logging.get_logger(__name__)

A   :					Union[str, 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"

A   :					Tuple	=      "\\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"

A   :					Tuple	=      "\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     _lowercase				(							datasets.Metric):


       """simple docstring"""




       def 		lowerCAmelCase  (     self							:    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 		lowerCAmelCase  (     self							:    Union[str, Any]      ,   __lowerCamelCase							:    str ):


             '''simple docstring'''



             if self.config_name == "default":
                   lowerCamelCase__	:       str					=						comet.load_from_checkpoint(comet.download_model("wmt20-comet-da" ) )
             else:
                   lowerCamelCase__	:       Optional[Any]					=						comet.load_from_checkpoint(comet.download_model(self.config_name ) )





       def 		lowerCAmelCase  (     self							:    Tuple      ,   __lowerCamelCase							:    str      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    Any=None      ,   __lowerCamelCase							:    Optional[int]=False ):


             '''simple docstring'''



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






import unittest

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

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
        OpenAIGPTConfig,
        OpenAIGPTDoubleHeadsModel,
        OpenAIGPTForSequenceClassification,
        OpenAIGPTLMHeadModel,
        OpenAIGPTModel,
    )







class     _lowercase				:


       """simple docstring"""




       def __init__(     self							:    Dict      ,   __lowerCamelCase							:    str      ,   __lowerCamelCase							:    Optional[int]=13      ,   __lowerCamelCase							:    List[str]=7      ,   __lowerCamelCase							:    Tuple=True      ,   __lowerCamelCase							:    Optional[int]=True      ,   __lowerCamelCase							:    List[str]=True      ,   __lowerCamelCase							:    Union[str, Any]=99      ,   __lowerCamelCase							:    List[Any]=32      ,   __lowerCamelCase							:    List[Any]=5      ,   __lowerCamelCase							:    Optional[Any]=4      ,   __lowerCamelCase							:    Optional[int]=37      ,   __lowerCamelCase							:    List[str]="gelu"      ,   __lowerCamelCase							:    List[str]=0.1      ,   __lowerCamelCase							:    int=0.1      ,   __lowerCamelCase							:    List[str]=512      ,   __lowerCamelCase							:    Optional[Any]=16      ,   __lowerCamelCase							:    Optional[Any]=2      ,   __lowerCamelCase							:    str=0.0_2      ,   __lowerCamelCase							:    List[str]=3      ,   __lowerCamelCase							:    Tuple=4      ,   __lowerCamelCase							:    Optional[int]=None      ,   ):


             '''simple docstring'''



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





       def 		lowerCAmelCase  (     self							:    List[Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[int]					=						ids_tensor([self.batch_size, self.seq_length]      ,   self.vocab_size )

             lowerCamelCase__	:       Optional[Any]					=						None
             if self.use_token_type_ids:
                   lowerCamelCase__	:       Optional[Any]					=						ids_tensor([self.batch_size, self.seq_length]      ,   self.type_vocab_size )

             lowerCamelCase__	:       Any					=						None
             lowerCamelCase__	:       str					=						None
             lowerCamelCase__	:       str					=						None
             if self.use_labels:
                   lowerCamelCase__	:       Dict					=						ids_tensor([self.batch_size]      ,   self.type_sequence_label_size )
                   lowerCamelCase__	:       Dict					=						ids_tensor([self.batch_size, self.seq_length]      ,   self.num_labels )
                   lowerCamelCase__	:       Dict					=						ids_tensor([self.batch_size]      ,   self.num_choices )

             lowerCamelCase__	:       Union[str, Any]					=						OpenAIGPTConfig(
                 vocab_size=self.vocab_size      ,   n_embd=self.hidden_size      ,   n_layer=self.num_hidden_layers      ,   n_head=self.num_attention_heads      ,   n_positions=self.max_position_embeddings      ,   pad_token_id=self.pad_token_id      ,   )

             lowerCamelCase__	:       Optional[int]					=						ids_tensor([self.num_hidden_layers, self.num_attention_heads]      ,   2 )

             return (
                 config,
                 input_ids,
                 head_mask,
                 token_type_ids,
                 sequence_labels,
                 token_labels,
                 choice_labels,
             )





       def 		lowerCAmelCase  (     self							:    str      ,   __lowerCamelCase							:    int      ,   __lowerCamelCase							:    int      ,   __lowerCamelCase							:    Optional[int]      ,   __lowerCamelCase							:    int      ,   *__lowerCamelCase							:    List[Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[int]					=						OpenAIGPTModel(config=__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()

             lowerCamelCase__	:       Tuple					=						model(__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   head_mask=__lowerCamelCase )
             lowerCamelCase__	:       str					=						model(__lowerCamelCase      ,   token_type_ids=__lowerCamelCase )
             lowerCamelCase__	:       Optional[int]					=						model(__lowerCamelCase )

             self.parent.assertEqual(result.last_hidden_state.shape      ,   (self.batch_size, self.seq_length, self.hidden_size) )





       def 		lowerCAmelCase  (     self							:    str      ,   __lowerCamelCase							:    List[Any]      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    List[str]      ,   __lowerCamelCase							:    Any      ,   *__lowerCamelCase							:    Optional[int] ):


             '''simple docstring'''



             lowerCamelCase__	:       Tuple					=						OpenAIGPTLMHeadModel(__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()

             lowerCamelCase__	:       List[str]					=						model(__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase )
             self.parent.assertEqual(result.loss.shape      ,   () )
             self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size) )





       def 		lowerCAmelCase  (     self							:    Dict      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    Optional[int]      ,   __lowerCamelCase							:    Optional[int]      ,   *__lowerCamelCase							:    Tuple ):


             '''simple docstring'''



             lowerCamelCase__	:       List[Any]					=						OpenAIGPTDoubleHeadsModel(__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()

             lowerCamelCase__	:       Optional[Any]					=						model(__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase )
             self.parent.assertEqual(result.loss.shape      ,   () )
             self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size) )





       def 		lowerCAmelCase  (     self							:    Tuple      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    List[Any]      ,   *__lowerCamelCase							:    Optional[int] ):


             '''simple docstring'''



             lowerCamelCase__	:       Dict					=						self.num_labels
             lowerCamelCase__	:       Tuple					=						OpenAIGPTForSequenceClassification(__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()

             lowerCamelCase__	:       Dict					=						ids_tensor([self.batch_size]      ,   self.type_sequence_label_size )
             lowerCamelCase__	:       List[str]					=						model(__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase )
             self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.num_labels) )





       def 		lowerCAmelCase  (     self							:    List[str] ):


             '''simple docstring'''



             lowerCamelCase__	:       str					=						self.prepare_config_and_inputs()
             (
                 (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     
             )	:       Any					=						config_and_inputs
             lowerCamelCase__	:       Union[str, Any]					=						{
                 "input_ids": input_ids,
                 "token_type_ids": token_type_ids,
                 "head_mask": head_mask,
             }

             return config, inputs_dict







@require_torch
class     _lowercase				(							lowercase__							,						lowercase__							,						lowercase__							,						unittest.TestCase):


       """simple docstring"""




       A__      		=    (
           (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
           if is_torch_available()
           else ()
       )
       A__      		=    (
           (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
       )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
       A__      		=    (
           {
               "feature-extraction": OpenAIGPTModel,
               "text-classification": OpenAIGPTForSequenceClassification,
               "text-generation": OpenAIGPTLMHeadModel,
               "zero-shot": OpenAIGPTForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )





       def 		lowerCAmelCase  (     self							:    List[str]      ,   __lowerCamelCase							:    str      ,   __lowerCamelCase							:    Tuple      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    List[Any]      ,   __lowerCamelCase							:    Union[str, Any] ):


             '''simple docstring'''



             if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
                   # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
                   # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
                   # tiny config could not be created.
                   return True

             return False





       def 		lowerCAmelCase  (     self							:    Union[str, Any]      ,   __lowerCamelCase							:    Optional[int]      ,   __lowerCamelCase							:    Tuple      ,   __lowerCamelCase							:    Tuple=False ):


             '''simple docstring'''



             lowerCamelCase__	:       Tuple					=						super()._prepare_for_class(__lowerCamelCase      ,   __lowerCamelCase      ,   return_labels=__lowerCamelCase )

             if return_labels:
                   if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
                         lowerCamelCase__	:       Optional[Any]					=						torch.zeros(
                             (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)      ,   dtype=torch.long      ,   device=__lowerCamelCase      ,   )
                         lowerCamelCase__	:       Tuple					=						inputs_dict["labels"]
                         lowerCamelCase__	:       Any					=						inputs_dict["labels"]
                         lowerCamelCase__	:       Any					=						torch.zeros(
                             (self.model_tester.batch_size, self.model_tester.num_choices)      ,   dtype=torch.long      ,   device=__lowerCamelCase      ,   )
                         lowerCamelCase__	:       Union[str, Any]					=						torch.zeros(
                             self.model_tester.batch_size      ,   dtype=torch.long      ,   device=__lowerCamelCase )
             return inputs_dict





       def 		lowerCAmelCase  (     self							:    List[Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       Tuple					=						OpenAIGPTModelTester(self )
             lowerCamelCase__	:       Union[str, Any]					=						ConfigTester(self      ,   config_class=__lowerCamelCase      ,   n_embd=37 )





       def 		lowerCAmelCase  (     self							:    int ):


             '''simple docstring'''



             self.config_tester.run_common_tests()





       def 		lowerCAmelCase  (     self							:    Dict ):


             '''simple docstring'''



             lowerCamelCase__	:       List[Any]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_openai_gpt_model(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    str ):


             '''simple docstring'''



             lowerCamelCase__	:       Union[str, Any]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_lm_head_model(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Dict ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[int]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_double_lm_head_model(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Optional[int] ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[int]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*__lowerCamelCase )





       @slow
       def 		lowerCAmelCase  (     self							:    List[str] ):


             '''simple docstring'''



             for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   lowerCamelCase__	:       Any					=						OpenAIGPTModel.from_pretrained(__lowerCamelCase )
                   self.assertIsNotNone(__lowerCamelCase )







@require_torch
class     _lowercase				(							unittest.TestCase):


       """simple docstring"""




       @slow
       def 		lowerCAmelCase  (     self							:    Any ):


             '''simple docstring'''



             lowerCamelCase__	:       List[Any]					=						OpenAIGPTLMHeadModel.from_pretrained("openai-gpt" )
             model.to(__lowerCamelCase )
             lowerCamelCase__	:       int					=						torch.tensor([[481, 4735, 544]]      ,   dtype=torch.long      ,   device=__lowerCamelCase )  # the president is
             lowerCamelCase__	:       Union[str, Any]					=						[
                 481,
                 4735,
                 544,
                 246,
                 963,
                 870,
                 762,
                 239,
                 244,
                 40477,
                 244,
                 249,
                 719,
                 881,
                 487,
                 544,
                 240,
                 244,
                 603,
                 481,
             ]  # the president is a very good man. " \n " i\'m sure he is, " said the

             lowerCamelCase__	:       int					=						model.generate(__lowerCamelCase      ,   do_sample=__lowerCamelCase )
             self.assertListEqual(output_ids[0].tolist()      ,   __lowerCamelCase )


