






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 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		List[str]    )       ->List[str]:




       """simple docstring"""


       a						    = '''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
       a						    = dict(zip(__UpperCAmelCase					,   range(len(__UpperCAmelCase    )    )    )    )

       a						    = {
           '''unk_token''': '''<unk>''',
           '''bos_token''': '''<s>''',
           '''eos_token''': '''</s>''',
       }
       a						    = {
           '''feature_size''': 1,
           '''padding_value''': 0.0,
           '''sampling_rate''': 16_000,
           '''return_attention_mask''': False,
           '''do_normalize''': True,
       }

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

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

       # load decoder from hub
       a						    = '''hf-internal-testing/ngram-beam-search-decoder'''







      def  __lowerCAmelCase   ( self :		str					,   **__UpperCAmelCase :		Optional[int]    )       ->Dict:




       """simple docstring"""


       a						    = self.add_kwargs_tokens_map.copy()
       kwargs.update(__UpperCAmelCase    )
       return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname					,   **__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		List[Any]					,   **__UpperCAmelCase :		Tuple    )       ->List[Any]:




       """simple docstring"""


       return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname					,   **__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		str					,   **__UpperCAmelCase :		Dict    )       ->Union[str, Any]:




       """simple docstring"""


       return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name					,   **__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		int    )       ->Tuple:




       """simple docstring"""


       shutil.rmtree(self.tmpdirname    )







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




       """simple docstring"""


       a						    = self.get_tokenizer()
       a						    = self.get_feature_extractor()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       processor.save_pretrained(self.tmpdirname    )
       a						    = WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname    )

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

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

       # 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					,   __UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Dict    )       ->str:




       """simple docstring"""


       a						    = 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
       a						    = 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  __lowerCAmelCase   ( self :		Any    )       ->List[str]:




       """simple docstring"""


       a						    = self.get_tokenizer()
       # add token to trigger raise
       tokenizer.add_tokens(['''xx''']    )
       with self.assertRaisesRegex(__UpperCAmelCase					,   '''include'''    ):
        WavaVecaProcessorWithLM(
            tokenizer=__UpperCAmelCase					,   feature_extractor=self.get_feature_extractor()					,   decoder=self.get_decoder()    )







      def  __lowerCAmelCase   ( self :		str    )       ->Any:




       """simple docstring"""


       a						    = self.get_feature_extractor()
       a						    = self.get_tokenizer()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       a						    = floats_list((3, 1_000)    )

       a						    = feature_extractor(__UpperCAmelCase					,   return_tensors='''np'''    )
       a						    = processor(__UpperCAmelCase					,   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  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Optional[Any]:




       """simple docstring"""


       a						    = self.get_feature_extractor()
       a						    = self.get_tokenizer()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       a						    = '''This is a test string'''

       a						    = processor(text=__UpperCAmelCase    )

       a						    = tokenizer(__UpperCAmelCase    )

       for key in encoded_tok.keys():
        self.assertListEqual(encoded_tok[key]					,   encoded_processor[key]    )







      def  __lowerCAmelCase   ( self :		int					,   __UpperCAmelCase :		Dict=(2, 10, 16)					,   __UpperCAmelCase :		Tuple=77    )       ->Dict:




       """simple docstring"""


       np.random.seed(__UpperCAmelCase    )
       return np.random.rand(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		int    )       ->Optional[Any]:




       """simple docstring"""


       a						    = self.get_feature_extractor()
       a						    = self.get_tokenizer()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       a						    = self._get_dummy_logits(shape=(10, 16)					,   seed=13    )

       a						    = processor.decode(__UpperCAmelCase    )

       a						    = decoder.decode_beams(__UpperCAmelCase    )[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  __lowerCAmelCase   ( self :		Dict					,   __UpperCAmelCase :		Union[str, Any]    )       ->Union[str, Any]:




       """simple docstring"""


       a						    = self.get_feature_extractor()
       a						    = self.get_tokenizer()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       a						    = 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:
        a						    = processor.batch_decode(__UpperCAmelCase    )
       else:
        with get_context(__UpperCAmelCase    ).Pool() as pool:
         a						    = processor.batch_decode(__UpperCAmelCase					,   __UpperCAmelCase    )

       a						    = list(__UpperCAmelCase    )

       with get_context('''fork'''    ).Pool() as p:
        a						    = decoder.decode_beams_batch(__UpperCAmelCase					,   __UpperCAmelCase    )

       a    ,	a    ,	a						    = [], [], []
       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(__UpperCAmelCase					,   decoded_processor.text    )
       self.assertListEqual(['''<s> <s> </s>''', '''<s> <s> <s>''']					,   decoded_processor.text    )
       self.assertListEqual(__UpperCAmelCase					,   decoded_processor.logit_score    )
       self.assertListEqual(__UpperCAmelCase					,   decoded_processor.lm_score    )







      def  __lowerCAmelCase   ( self :		List[Any]    )       ->int:




       """simple docstring"""


       a						    = self.get_feature_extractor()
       a						    = self.get_tokenizer()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       a						    = self._get_dummy_logits()

       a						    = 15
       a						    = -20.0
       a						    = -4.0

       a						    = processor.batch_decode(
           __UpperCAmelCase					,   beam_width=__UpperCAmelCase					,   beam_prune_logp=__UpperCAmelCase					,   token_min_logp=__UpperCAmelCase					,   )
       a						    = decoded_processor_out.text

       a						    = list(__UpperCAmelCase    )

       with get_context('''fork'''    ).Pool() as pool:
        a						    = decoder.decode_beams_batch(
            __UpperCAmelCase					,   __UpperCAmelCase					,   beam_width=__UpperCAmelCase					,   beam_prune_logp=__UpperCAmelCase					,   token_min_logp=__UpperCAmelCase					,   )

       a						    = [d[0][0] for d in decoded_decoder_out]
       a						    = [d[0][2] for d in decoded_decoder_out]
       a						    = [d[0][3] for d in decoded_decoder_out]

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

       self.assertTrue(np.array_equal(__UpperCAmelCase					,   decoded_processor_out.logit_score    )    )
       self.assertTrue(np.allclose([-20.054, -18.447]					,   __UpperCAmelCase					,   atol=1e-3    )    )

       self.assertTrue(np.array_equal(__UpperCAmelCase					,   decoded_processor_out.lm_score    )    )
       self.assertTrue(np.allclose([-15.554, -13.9474]					,   __UpperCAmelCase					,   atol=1e-3    )    )







      def  __lowerCAmelCase   ( self :		int    )       ->str:




       """simple docstring"""


       a						    = self.get_feature_extractor()
       a						    = self.get_tokenizer()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       a						    = self._get_dummy_logits()

       a						    = 2.0
       a						    = 5.0
       a						    = -20.0
       a						    = True

       a						    = processor.batch_decode(
           __UpperCAmelCase					,   alpha=__UpperCAmelCase					,   beta=__UpperCAmelCase					,   unk_score_offset=__UpperCAmelCase					,   lm_score_boundary=__UpperCAmelCase					,   )
       a						    = decoded_processor_out.text

       a						    = list(__UpperCAmelCase    )
       decoder.reset_params(
           alpha=__UpperCAmelCase					,   beta=__UpperCAmelCase					,   unk_score_offset=__UpperCAmelCase					,   lm_score_boundary=__UpperCAmelCase					,   )

       with get_context('''fork'''    ).Pool() as pool:
        a						    = decoder.decode_beams_batch(
            __UpperCAmelCase					,   __UpperCAmelCase					,   )

       a						    = [d[0][0] for d in decoded_decoder_out]

       self.assertListEqual(__UpperCAmelCase					,   __UpperCAmelCase    )
       self.assertListEqual(['''<s> </s> <s> </s> </s>''', '''</s> </s> <s> </s> </s>''']					,   __UpperCAmelCase    )
       a						    = 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					,   __UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		str    )       ->Any:




       """simple docstring"""


       a						    = WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''    )

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

       a						    = os.listdir(__UpperCAmelCase    )
       a						    = ['''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(__UpperCAmelCase					,   __UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Optional[Any]:




       """simple docstring"""


       a						    = snapshot_download('''hf-internal-testing/processor_with_lm'''    )

       a						    = WavaVecaProcessorWithLM.from_pretrained(__UpperCAmelCase    )

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

       a						    = os.listdir(__UpperCAmelCase    )
       a						    = os.listdir(__UpperCAmelCase    )

       local_decoder_files.sort()
       expected_decoder_files.sort()

       # test that both decoder form hub and local files in cache are the same
       self.assertListEqual(__UpperCAmelCase					,   __UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Dict    )       ->Dict:




       """simple docstring"""


       a						    = WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''    )
       a						    = AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm'''    )

       a						    = floats_list((3, 1_000)    )

       a						    = processor_wavaveca(__UpperCAmelCase					,   return_tensors='''np'''    )
       a						    = processor_auto(__UpperCAmelCase					,   return_tensors='''np'''    )

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

       a						    = self._get_dummy_logits()

       a						    = processor_wavaveca.batch_decode(__UpperCAmelCase    )
       a						    = processor_auto.batch_decode(__UpperCAmelCase    )

       self.assertListEqual(decoded_wavaveca.text					,   decoded_auto.text    )







      def  __lowerCAmelCase   ( self :		Optional[int]    )       ->Optional[Any]:




       """simple docstring"""


       a						    = self.get_feature_extractor()
       a						    = self.get_tokenizer()
       a						    = self.get_decoder()

       a						    = WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase					,   feature_extractor=__UpperCAmelCase					,   decoder=__UpperCAmelCase    )

       self.assertListEqual(
           processor.model_input_names					,   feature_extractor.model_input_names					,   msg='''`processor` and `feature_extractor` model input names do not match'''					,   )







      @staticmethod
      def  __lowerCAmelCase   ( __UpperCAmelCase :		int					,   __UpperCAmelCase :		List[str]    )       ->Optional[int]:




       """simple docstring"""


       a						    = [d[key] for d in offsets]
       return retrieved_list







      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Tuple:




       """simple docstring"""


       a						    = WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''    )
       a						    = self._get_dummy_logits()[0]

       a						    = processor.decode(__UpperCAmelCase					,   output_word_offsets=__UpperCAmelCase    )
       # 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(__UpperCAmelCase					,   __UpperCAmelCase    )    )

       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  __lowerCAmelCase   ( self :		List[str]    )       ->Optional[int]:




       """simple docstring"""


       a						    = WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''    )
       a						    = self._get_dummy_logits()

       a						    = processor.batch_decode(__UpperCAmelCase					,   output_word_offsets=__UpperCAmelCase    )

       # 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(__UpperCAmelCase					,   __UpperCAmelCase    )    )

       self.assertListEqual(
           [''' '''.join(self.get_from_offsets(__UpperCAmelCase					,   '''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  __lowerCAmelCase   ( self :		List[str]    )       ->Any:




       """simple docstring"""


       import torch

       a						    = load_dataset('''common_voice'''					,   '''en'''					,   split='''train'''					,   streaming=__UpperCAmelCase    )
       a						    = ds.cast_column('''audio'''					,   datasets.Audio(sampling_rate=16_000    )    )
       a						    = iter(__UpperCAmelCase    )
       a						    = next(__UpperCAmelCase    )

       a						    = AutoProcessor.from_pretrained('''patrickvonplaten/wav2vec2-base-100h-with-lm'''    )
       a						    = 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
       a						    = processor(sample['''audio''']['''array''']					,   return_tensors='''pt'''    ).input_values

       with torch.no_grad():
        a						    = model(__UpperCAmelCase    ).logits.cpu().numpy()

       a						    = processor.decode(logits[0]					,   output_word_offsets=__UpperCAmelCase    )

       a						    = model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
       a						    = [
           {
               '''start_time''': d['''start_offset'''] * time_offset,
               '''end_time''': d['''end_offset'''] * time_offset,
               '''word''': d['''word'''],
           }
           for d in output['''word_offsets''']
       ]

       a						    = '''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(__UpperCAmelCase					,   '''word'''    )    )					,   __UpperCAmelCase    )
       self.assertEqual(''' '''.join(self.get_from_offsets(__UpperCAmelCase					,   '''word'''    )    )					,   output.text    )

       # output times
       a						    = torch.tensor(self.get_from_offsets(__UpperCAmelCase					,   '''start_time'''    )    )
       a						    = torch.tensor(self.get_from_offsets(__UpperCAmelCase					,   '''end_time'''    )    )

       # fmt: off
       a						    = torch.tensor([1.4199, 1.6599, 2.2599, 3.0, 3.24, 3.5999, 3.7999, 4.0999, 4.26, 4.94, 5.28, 5.6599, 5.78, 5.94, 6.32, 6.5399, 6.6599]    )
       a						    = torch.tensor([1.5399, 1.8999, 2.9, 3.16, 3.5399, 3.72, 4.0199, 4.1799, 4.76, 5.1599, 5.5599, 5.6999, 5.86, 6.1999, 6.38, 6.6199, 6.94]    )
       # fmt: on

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










import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

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


UpperCAmelCase__							  =						logging.get_logger(__name__)

UpperCAmelCase__							  =						"▁"

UpperCAmelCase__							  =						{"vocab_file": "sentencepiece.bpe.model"}

UpperCAmelCase__							  =						{
    "vocab_file": {
        "facebook/nllb-200-distilled-600M": (
            "https://huggingface.co/facebook/nllb-200-distilled-600M/blob/main/sentencepiece.bpe.model"
        ),
    }
}

UpperCAmelCase__							  =						{
    "facebook/nllb-200-distilled-600M": 1024,
}

# fmt: off
UpperCAmelCase__							  =						["ace_Arab", "ace_Latn", "acm_Arab", "acq_Arab", "aeb_Arab", "afr_Latn", "ajp_Arab", "aka_Latn", "amh_Ethi", "apc_Arab", "arb_Arab", "ars_Arab", "ary_Arab", "arz_Arab", "asm_Beng", "ast_Latn", "awa_Deva", "ayr_Latn", "azb_Arab", "azj_Latn", "bak_Cyrl", "bam_Latn", "ban_Latn", "bel_Cyrl", "bem_Latn", "ben_Beng", "bho_Deva", "bjn_Arab", "bjn_Latn", "bod_Tibt", "bos_Latn", "bug_Latn", "bul_Cyrl", "cat_Latn", "ceb_Latn", "ces_Latn", "cjk_Latn", "ckb_Arab", "crh_Latn", "cym_Latn", "dan_Latn", "deu_Latn", "dik_Latn", "dyu_Latn", "dzo_Tibt", "ell_Grek", "eng_Latn", "epo_Latn", "est_Latn", "eus_Latn", "ewe_Latn", "fao_Latn", "pes_Arab", "fij_Latn", "fin_Latn", "fon_Latn", "fra_Latn", "fur_Latn", "fuv_Latn", "gla_Latn", "gle_Latn", "glg_Latn", "grn_Latn", "guj_Gujr", "hat_Latn", "hau_Latn", "heb_Hebr", "hin_Deva", "hne_Deva", "hrv_Latn", "hun_Latn", "hye_Armn", "ibo_Latn", "ilo_Latn", "ind_Latn", "isl_Latn", "ita_Latn", "jav_Latn", "jpn_Jpan", "kab_Latn", "kac_Latn", "kam_Latn", "kan_Knda", "kas_Arab", "kas_Deva", "kat_Geor", "knc_Arab", "knc_Latn", "kaz_Cyrl", "kbp_Latn", "kea_Latn", "khm_Khmr", "kik_Latn", "kin_Latn", "kir_Cyrl", "kmb_Latn", "kon_Latn", "kor_Hang", "kmr_Latn", "lao_Laoo", "lvs_Latn", "lij_Latn", "lim_Latn", "lin_Latn", "lit_Latn", "lmo_Latn", "ltg_Latn", "ltz_Latn", "lua_Latn", "lug_Latn", "luo_Latn", "lus_Latn", "mag_Deva", "mai_Deva", "mal_Mlym", "mar_Deva", "min_Latn", "mkd_Cyrl", "plt_Latn", "mlt_Latn", "mni_Beng", "khk_Cyrl", "mos_Latn", "mri_Latn", "zsm_Latn", "mya_Mymr", "nld_Latn", "nno_Latn", "nob_Latn", "npi_Deva", "nso_Latn", "nus_Latn", "nya_Latn", "oci_Latn", "gaz_Latn", "ory_Orya", "pag_Latn", "pan_Guru", "pap_Latn", "pol_Latn", "por_Latn", "prs_Arab", "pbt_Arab", "quy_Latn", "ron_Latn", "run_Latn", "rus_Cyrl", "sag_Latn", "san_Deva", "sat_Beng", "scn_Latn", "shn_Mymr", "sin_Sinh", "slk_Latn", "slv_Latn", "smo_Latn", "sna_Latn", "snd_Arab", "som_Latn", "sot_Latn", "spa_Latn", "als_Latn", "srd_Latn", "srp_Cyrl", "ssw_Latn", "sun_Latn", "swe_Latn", "swh_Latn", "szl_Latn", "tam_Taml", "tat_Cyrl", "tel_Telu", "tgk_Cyrl", "tgl_Latn", "tha_Thai", "tir_Ethi", "taq_Latn", "taq_Tfng", "tpi_Latn", "tsn_Latn", "tso_Latn", "tuk_Latn", "tum_Latn", "tur_Latn", "twi_Latn", "tzm_Tfng", "uig_Arab", "ukr_Cyrl", "umb_Latn", "urd_Arab", "uzn_Latn", "vec_Latn", "vie_Latn", "war_Latn", "wol_Latn", "xho_Latn", "ydd_Hebr", "yor_Latn", "yue_Hant", "zho_Hans", "zho_Hant", "zul_Latn"]






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




      __snake_case	    =    VOCAB_FILES_NAMES
      __snake_case	    =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __snake_case	    =    PRETRAINED_VOCAB_FILES_MAP
      __snake_case	    =    ['''input_ids''', '''attention_mask''']

      __snake_case	    =    []
      __snake_case	    =    []







      def __init__( self :		int					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[int]="<s>"					,   __UpperCAmelCase :		Union[str, Any]="</s>"					,   __UpperCAmelCase :		List[Any]="</s>"					,   __UpperCAmelCase :		List[Any]="<s>"					,   __UpperCAmelCase :		Optional[int]="<unk>"					,   __UpperCAmelCase :		Optional[int]="<pad>"					,   __UpperCAmelCase :		Tuple="<mask>"					,   __UpperCAmelCase :		Any=None					,   __UpperCAmelCase :		Dict=None					,   __UpperCAmelCase :		int=None					,   __UpperCAmelCase :		Optional[Dict[str, Any]] = None					,   __UpperCAmelCase :		List[Any]=None					,   __UpperCAmelCase :		int=False					,   **__UpperCAmelCase :		Dict					,   )       ->List[str]:




       """simple docstring"""


       a						    = AddedToken(__UpperCAmelCase					,   lstrip=__UpperCAmelCase					,   rstrip=__UpperCAmelCase    ) if isinstance(__UpperCAmelCase					,   __UpperCAmelCase    ) else mask_token

       a						    = {} if sp_model_kwargs is None else sp_model_kwargs
       a						    = legacy_behaviour

       super().__init__(
           bos_token=__UpperCAmelCase					,   eos_token=__UpperCAmelCase					,   unk_token=__UpperCAmelCase					,   sep_token=__UpperCAmelCase					,   cls_token=__UpperCAmelCase					,   pad_token=__UpperCAmelCase					,   mask_token=__UpperCAmelCase					,   tokenizer_file=__UpperCAmelCase					,   src_lang=__UpperCAmelCase					,   tgt_lang=__UpperCAmelCase					,   additional_special_tokens=__UpperCAmelCase					,   sp_model_kwargs=self.sp_model_kwargs					,   legacy_behaviour=__UpperCAmelCase					,   **__UpperCAmelCase					,   )

       a						    = spm.SentencePieceProcessor(**self.sp_model_kwargs    )
       self.sp_model.Load(str(__UpperCAmelCase    )    )
       a						    = vocab_file
       # Original fairseq vocab and spm vocab must be "aligned":
       # Vocab    |    0    |    1    |   2    |    3    |  4   |  5   |  6   |   7  |   8  |  9
       # -------- | ------- | ------- | ------ | ------- | ---- | ---- | ---- | ---- | ---- | ----
       # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | 'an' | '▁n' | '▁m' | '▁t' | '▁k' | '▁a'
       # spm      | '<unk>' | '<s>'   | '</s>' | 'an'    | '▁n' | '▁m' | '▁t' | '▁k' | '▁a' | '▁s'

       # Mimic fairseq token-to-id alignment for the first 4 token
       a						    = {'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

       # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
       a						    = 1

       a						    = len(self.sp_model    )
       a						    = {
           code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(__UpperCAmelCase    )
       }
       a						    = {v: k for k, v in self.lang_code_to_id.items()}
       a						    = len(self.sp_model    ) + len(self.lang_code_to_id    ) + self.fairseq_offset

       self.fairseq_tokens_to_ids.update(self.lang_code_to_id    )
       a						    = {v: k for k, v in self.fairseq_tokens_to_ids.items()}
       a						    = list(self.lang_code_to_id.keys()    )

       if additional_special_tokens is not None:
        # Only add those special tokens if they are not already there.
        self._additional_special_tokens.extend(
            [t for t in additional_special_tokens if t not in self._additional_special_tokens]    )

       a						    = src_lang if src_lang is not None else '''eng_Latn'''
       a						    = self.lang_code_to_id[self._src_lang]
       a						    = tgt_lang
       self.set_src_lang_special_tokens(self._src_lang    )







      def __getstate__( self :		str    )       ->Any:




       """simple docstring"""


       a						    = self.__dict__.copy()
       a						    = None
       a						    = self.sp_model.serialized_model_proto()
       return state







      def __setstate__( self :		List[str]					,   __UpperCAmelCase :		Optional[Any]    )       ->int:




       """simple docstring"""


       a						    = d

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

       a						    = spm.SentencePieceProcessor(**self.sp_model_kwargs    )
       self.sp_model.LoadFromSerializedProto(self.sp_model_proto    )







      @property
      def  __lowerCAmelCase   ( self :		Tuple    )       ->Optional[Any]:




       """simple docstring"""


       return len(self.sp_model    ) + len(self.lang_code_to_id    ) + self.fairseq_offset + 1  # Plus 1 for the mask token







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->str:




       """simple docstring"""


       return self._src_lang







      @src_lang.setter
      def  __lowerCAmelCase   ( self :		List[Any]					,   __UpperCAmelCase :		str    )       ->None:




       """simple docstring"""


       a						    = new_src_lang
       self.set_src_lang_special_tokens(self._src_lang    )







      def  __lowerCAmelCase   ( self :		Any					,   __UpperCAmelCase :		List[int]					,   __UpperCAmelCase :		Optional[List[int]] = None					,   __UpperCAmelCase :		bool = False    )       ->List[int]:




       """simple docstring"""


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

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







      def  __lowerCAmelCase   ( self :		Any					,   __UpperCAmelCase :		List[int]					,   __UpperCAmelCase :		Optional[List[int]] = None    )       ->List[int]:




       """simple docstring"""


       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  __lowerCAmelCase   ( self :		int					,   __UpperCAmelCase :		List[int]					,   __UpperCAmelCase :		Optional[List[int]] = None    )       ->List[int]:




       """simple docstring"""


       a						    = [self.sep_token_id]
       a						    = [self.cls_token_id]

       if token_ids_a is None:
        return len(cls + token_ids_a + sep    ) * [0]
       return len(cls + token_ids_a + sep + sep + token_ids_a + sep    ) * [0]







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		Optional[Any]					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[str]					,   __UpperCAmelCase :		Optional[str]					,   **__UpperCAmelCase :		Tuple    )       ->Tuple:




       """simple docstring"""


       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(__UpperCAmelCase					,   add_special_tokens=__UpperCAmelCase					,   return_tensors=__UpperCAmelCase					,   **__UpperCAmelCase    )
       a						    = self.convert_tokens_to_ids(__UpperCAmelCase    )
       a						    = tgt_lang_id
       return inputs







      def  __lowerCAmelCase   ( self :		Optional[int]    )       ->Union[str, Any]:




       """simple docstring"""


       a						    = {self.convert_ids_to_tokens(__UpperCAmelCase    ): i for i in range(self.vocab_size    )}
       vocab.update(self.added_tokens_encoder    )
       return vocab







      def  __lowerCAmelCase   ( self :		List[Any]					,   __UpperCAmelCase :		str    )       ->List[str]:




       """simple docstring"""


       return self.sp_model.encode(__UpperCAmelCase					,   out_type=__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Optional[int]					,   __UpperCAmelCase :		str    )       ->List[Any]:




       """simple docstring"""


       if token in self.fairseq_tokens_to_ids:
        return self.fairseq_tokens_to_ids[token]
       a						    = self.sp_model.PieceToId(__UpperCAmelCase    )

       # Need to return unknown token if the SP model returned 0
       return spm_id + self.fairseq_offset if spm_id else self.unk_token_id







      def  __lowerCAmelCase   ( self :		Union[str, Any]					,   __UpperCAmelCase :		int    )       ->Optional[Any]:




       """simple docstring"""


       if index in self.fairseq_ids_to_tokens:
        return self.fairseq_ids_to_tokens[index]
       return self.sp_model.IdToPiece(index - self.fairseq_offset    )







      def  __lowerCAmelCase   ( self :		Tuple					,   __UpperCAmelCase :		List[str]    )       ->List[Any]:




       """simple docstring"""


       a						    = ''''''.join(__UpperCAmelCase    ).replace(__UpperCAmelCase					,   ''' '''    ).strip()
       return out_string







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[str] = None    )       ->Tuple[str]:




       """simple docstring"""


       if not os.path.isdir(__UpperCAmelCase    ):
        logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""    )
        return
       a						    = os.path.join(
           __UpperCAmelCase					,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']    )

       if os.path.abspath(self.vocab_file    ) != os.path.abspath(__UpperCAmelCase    ) and os.path.isfile(self.vocab_file    ):
        copyfile(self.vocab_file					,   __UpperCAmelCase    )
       elif not os.path.isfile(self.vocab_file    ):
        with open(__UpperCAmelCase					,   '''wb'''    ) as fi:
         a						    = self.sp_model.serialized_model_proto()
         fi.write(__UpperCAmelCase    )

       return (out_vocab_file,)







      def  __lowerCAmelCase   ( self :		List[Any]					,   __UpperCAmelCase :		List[str]					,   __UpperCAmelCase :		str = "eng_Latn"					,   __UpperCAmelCase :		Optional[List[str]] = None					,   __UpperCAmelCase :		str = "fra_Latn"					,   **__UpperCAmelCase :		Optional[int]					,   )       ->BatchEncoding:




       """simple docstring"""


       a						    = src_lang
       a						    = tgt_lang
       return super().prepare_seqaseq_batch(__UpperCAmelCase					,   __UpperCAmelCase					,   **__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Dict    )       ->List[str]:




       """simple docstring"""


       return self.set_src_lang_special_tokens(self.src_lang    )







      def  __lowerCAmelCase   ( self :		Any    )       ->Any:




       """simple docstring"""


       return self.set_tgt_lang_special_tokens(self.tgt_lang    )







      def  __lowerCAmelCase   ( self :		Any					,   __UpperCAmelCase :		Tuple    )       ->None:




       """simple docstring"""


       a						    = self.lang_code_to_id[src_lang]
       if self.legacy_behaviour:
        a						    = []
        a						    = [self.eos_token_id, self.cur_lang_code]
       else:
        a						    = [self.cur_lang_code]
        a						    = [self.eos_token_id]







      def  __lowerCAmelCase   ( self :		List[str]					,   __UpperCAmelCase :		str    )       ->None:




       """simple docstring"""


       a						    = self.lang_code_to_id[lang]
       if self.legacy_behaviour:
        a						    = []
        a						    = [self.eos_token_id, self.cur_lang_code]
       else:
        a						    = [self.cur_lang_code]
        a						    = [self.eos_token_id]


