


import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import BatchEncoding, PreTrainedTokenizer
from ...utils import logging


SCREAMING_SNAKE_CASE							    =							logging.get_logger(__name__)

SCREAMING_SNAKE_CASE							    =							'▁'

SCREAMING_SNAKE_CASE							    =							{
    'vocab_file': 'vocab.json',
    'spm_file': 'sentencepiece.bpe.model',
    'tokenizer_config_file': 'tokenizer_config.json',
}

SCREAMING_SNAKE_CASE							    =							{
    'vocab_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json',
    },
    'spm_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model',
    },
    'tokenizer_config_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json',
    },
}

SCREAMING_SNAKE_CASE							    =							{
    'facebook/m2m100_418M': 1_0_2_4,
}

# fmt: off
SCREAMING_SNAKE_CASE							    =							{
    'm2m100': ['af', 'am', 'ar', 'ast', 'az', 'ba', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'ceb', 'cs', 'cy', 'da', 'de', 'el', 'en', 'es', 'et', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gu', 'ha', 'he', 'hi', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'ilo', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'lb', 'lg', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'ne', 'nl', 'no', 'ns', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'ro', 'ru', 'sd', 'si', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'th', 'tl', 'tn', 'tr', 'uk', 'ur', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh', 'zu'],
    'wmt21': ['en', 'ha', 'is', 'ja', 'cs', 'ru', 'zh', 'de']
}


class     __UpperCAmelCase							(       __A							):


     """simple docstring"""
     _lowerCamelCase   =      VOCAB_FILES_NAMES
     _lowerCamelCase   =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     _lowerCamelCase   =      PRETRAINED_VOCAB_FILES_MAP
     _lowerCamelCase   =      ["""input_ids""", """attention_mask"""]

     _lowerCamelCase   =      []
     _lowerCamelCase   =      []


     def __init__(	self  ,       __A  ,       __A  ,       __A=None  ,       __A=None  ,       __A="<s>"  ,       __A="</s>"  ,       __A="</s>"  ,       __A="<pad>"  ,       __A="<unk>"  ,       __A="m2m100"  ,       __A = None  ,       __A=8  ,       **__A  ,       ):
        __a					=  {} if sp_model_kwargs is None else sp_model_kwargs

        __a					=  language_codes
        __a					=  FAIRSEQ_LANGUAGE_CODES[language_codes]
        __a					=  {lang_code: f'''__{lang_code}__''' for lang_code in fairseq_language_code}

        __a					=  kwargs.get("""additional_special_tokens"""  ,       []    )
        kwargs["additional_special_tokens"] += [
            self.get_lang_token(__A    )
            for lang_code in fairseq_language_code
            if self.get_lang_token(__A    ) not in kwargs["additional_special_tokens"]
        ]

        super().__init__(
            src_lang=__A  ,       tgt_lang=__A  ,       bos_token=__A  ,       eos_token=__A  ,       sep_token=__A  ,       unk_token=__A  ,       pad_token=__A  ,       language_codes=__A  ,       sp_model_kwargs=self.sp_model_kwargs  ,       num_madeup_words=__A  ,       **__A  ,       )

        __a					=  vocab_file
        __a					=  load_json(__A    )
        __a					=  {v: k for k, v in self.encoder.items()}
        __a					=  spm_file
        __a					=  load_spm(__A  ,       self.sp_model_kwargs    )

        __a					=  len(self.encoder    )

        __a					=  {
            self.get_lang_token(__A    ): self.encoder_size + i for i, lang_code in enumerate(__A    )
        }
        __a					=  {lang_code: self.encoder_size + i for i, lang_code in enumerate(__A    )}
        __a					=  {v: k for k, v in self.lang_token_to_id.items()}

        __a					=  src_lang if src_lang is not None else """en"""
        __a					=  tgt_lang
        __a					=  self.get_lang_id(self._src_lang    )
        self.set_src_lang_special_tokens(self._src_lang    )

        __a					=  num_madeup_words


     @property
     def 	snake_case_  (	self    ):
        return len(self.encoder    ) + len(self.lang_token_to_id    )


     @property
     def 	snake_case_  (	self    ):
        return self._src_lang


     @src_lang.setter
     def 	snake_case_  (	self  ,       __A    ):
        __a					=  new_src_lang
        self.set_src_lang_special_tokens(self._src_lang    )


     def 	snake_case_  (	self  ,       __A    ):
        return self.sp_model.encode(__A  ,       out_type=__A    )


     def 	snake_case_  (	self  ,       __A    ):
        if token in self.lang_token_to_id:
           return self.lang_token_to_id[token]
        return self.encoder.get(__A  ,       self.encoder[self.unk_token]    )


     def 	snake_case_  (	self  ,       __A    ):
        if index in self.id_to_lang_token:
           return self.id_to_lang_token[index]
        return self.decoder.get(__A  ,       self.unk_token    )


     def 	snake_case_  (	self  ,       __A    ):
        __a					=  []
        __a					=  """"""
        for token in tokens:
           # make sure that special tokens are not decoded using sentencepiece model
           if token in self.all_special_tokens:
              out_string += self.sp_model.decode(__A    ) + token
              __a					=  []
           else:
              current_sub_tokens.append(__A    )
        out_string += self.sp_model.decode(__A    )
        return out_string.strip()


     def 	snake_case_  (	self  ,       __A  ,       __A = None  ,       __A = False    ):

        if already_has_special_tokens:
           return super().get_special_tokens_mask(
               token_ids_a=__A  ,       token_ids_a=__A  ,       already_has_special_tokens=__A    )

        __a					=  [1] * len(self.prefix_tokens    )
        __a					=  [1] * len(self.suffix_tokens    )
        if token_ids_a is None:
           return prefix_ones + ([0] * len(__A    )) + suffix_ones
        return prefix_ones + ([0] * len(__A    )) + ([0] * len(__A    )) + suffix_ones


     def 	snake_case_  (	self  ,       __A  ,       __A = None    ):
        if token_ids_a is None:
           return self.prefix_tokens + token_ids_a + self.suffix_tokens
        # We don't expect to process pairs, but leave the pair logic for API consistency
        return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens


     def 	snake_case_  (	self    ):
        __a					=  {self.convert_ids_to_tokens(__A    ): i for i in range(self.vocab_size    )}
        vocab.update(self.added_tokens_encoder    )
        return vocab


     def __getstate__(	self    ):
        __a					=  self.__dict__.copy()
        __a					=  None
        return state


     def __setstate__(	self  ,       __A    ):
        __a					=  d

        # for backward compatibility
        if not hasattr(self  ,       """sp_model_kwargs"""    ):
           __a					=  {}

        __a					=  load_spm(self.spm_file  ,       self.sp_model_kwargs    )


     def 	snake_case_  (	self  ,       __A  ,       __A = None    ):
        __a					=  Path(__A    )
        if not save_dir.is_dir():
           raise OSError(f'''{save_directory} should be a directory'''    )
        __a					=  save_dir / (
            (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
        )
        __a					=  save_dir / (
            (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
        )

        save_json(self.encoder  ,       __A    )

        if os.path.abspath(self.spm_file    ) != os.path.abspath(__A    ) and os.path.isfile(self.spm_file    ):
           copyfile(self.spm_file  ,       __A    )
        elif not os.path.isfile(self.spm_file    ):
           with open(__A  ,       """wb"""    ) as fi:
              __a					=  self.sp_model.serialized_model_proto()
              fi.write(__A    )

        return (str(__A    ), str(__A    ))


     def 	snake_case_  (	self  ,       __A  ,       __A = "en"  ,       __A = None  ,       __A = "ro"  ,       **__A  ,       ):
        __a					=  src_lang
        __a					=  tgt_lang
        self.set_src_lang_special_tokens(self.src_lang    )
        return super().prepare_seqaseq_batch(__A  ,       __A  ,       **__A    )


     def 	snake_case_  (	self  ,       __A  ,       __A  ,       __A  ,       **__A    ):
        if src_lang is None or tgt_lang is None:
           raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""    )
        __a					=  src_lang
        __a					=  self(__A  ,       add_special_tokens=__A  ,       **__A    )
        __a					=  self.get_lang_id(__A    )
        __a					=  tgt_lang_id
        return inputs


     def 	snake_case_  (	self    ):
        self.set_src_lang_special_tokens(self.src_lang    )


     def 	snake_case_  (	self    ):
        self.set_tgt_lang_special_tokens(self.tgt_lang    )


     def 	snake_case_  (	self  ,       __A    ):
        __a					=  self.get_lang_token(__A    )
        __a					=  self.lang_token_to_id[lang_token]
        __a					=  [self.cur_lang_id]
        __a					=  [self.eos_token_id]


     def 	snake_case_  (	self  ,       __A    ):
        __a					=  self.get_lang_token(__A    )
        __a					=  self.lang_token_to_id[lang_token]
        __a					=  [self.cur_lang_id]
        __a					=  [self.eos_token_id]


     def 	snake_case_  (	self  ,       __A    ):
        return self.lang_code_to_token[lang]


     def 	snake_case_  (	self  ,       __A    ):
        __a					=  self.get_lang_token(__A    )
        return self.lang_token_to_id[lang_token]







def  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):
   __a					=  sentencepiece.SentencePieceProcessor(**lowerCAmelCase__   )
   spm.Load(str(lowerCAmelCase__   )   )
   return spm




def  a   (lowerCAmelCase__   ):
   with open(lowerCAmelCase__		,				"""r"""   ) as f:
      return json.load(lowerCAmelCase__   )




def  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):
   with open(lowerCAmelCase__		,				"""w"""   ) as f:
      json.dump(lowerCAmelCase__		,				lowerCAmelCase__		,				indent=2   )







import comet  # From: unbabel-comet
import torch

import datasets


SCREAMING_SNAKE_CASE							    =							datasets.logging.get_logger(__name__)

SCREAMING_SNAKE_CASE							    =							'\\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'

SCREAMING_SNAKE_CASE							    =							'\\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'

SCREAMING_SNAKE_CASE							    =							'\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     __UpperCAmelCase							(       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":
           __a					=  comet.load_from_checkpoint(comet.download_model("""wmt20-comet-da"""    )    )
        else:
           __a					=  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:
           __a					=  1 if torch.cuda.is_available() else 0
        __a					=  {"""src""": sources, """mt""": predictions, """ref""": references}
        __a					=  [dict(zip(__A  ,       __A    )    ) for t in zip(*data.values()    )]
        __a	,    __a					=  self.scorer.predict(__A  ,       gpus=__A  ,       progress_bar=__A    )
        return {"mean_score": mean_score, "scores": scores}



