





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):







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					List[Any]										="| <pad> <unk> <s> </s> a b c d e f g h i j k".split()
        a__	:					Optional[Any]										=dict(zip(lowerCAmelCase__		,      range(len(lowerCAmelCase__							)							)							)							)

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

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

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

        # load decoder from hub
        a__	:					str										="hf-internal-testing/ngram-beam-search-decoder"







      def    _lowercase  (    self		,      **lowerCAmelCase__							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					List[str]										=self.add_kwargs_tokens_map.copy()
        kwargs.update(lowerCAmelCase__							)
        return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname		,      **lowerCAmelCase__							)







      def    _lowercase  (    self		,      **lowerCAmelCase__							)   ->						Optional[int]:



        '''simple docstring'''




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







      def    _lowercase  (    self		,      **lowerCAmelCase__							)   ->						List[str]:



        '''simple docstring'''




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







      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        shutil.rmtree(self.tmpdirname							)







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					Any										=self.get_tokenizer()
        a__	:					str										=self.get_feature_extractor()
        a__	:					Tuple										=self.get_decoder()

        a__	:					Optional[Any]										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

        processor.save_pretrained(self.tmpdirname							)
        a__	:					List[Any]										=WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname							)

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

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

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







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



        '''simple docstring'''




        a__	:					Optional[int]										=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__	:					Tuple										=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    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Optional[Any]										=self.get_tokenizer()
        # add token to trigger raise
        tokenizer.add_tokens(["xx"]							)
        with self.assertRaisesRegex(lowerCAmelCase__		,      "include"							):
          WavaVecaProcessorWithLM(
              tokenizer=lowerCAmelCase__		,      feature_extractor=self.get_feature_extractor()		,      decoder=self.get_decoder()							)







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



        '''simple docstring'''




        a__	:					List[Any]										=self.get_feature_extractor()
        a__	:					List[Any]										=self.get_tokenizer()
        a__	:					List[str]										=self.get_decoder()

        a__	:					Optional[Any]										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

        a__	:					Optional[int]										=floats_list((3, 1_0_0_0)							)

        a__	:					int										=feature_extractor(lowerCAmelCase__		,      return_tensors="np"							)
        a__	:					Union[str, Any]										=processor(lowerCAmelCase__		,      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    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					str										=self.get_feature_extractor()
        a__	:					int										=self.get_tokenizer()
        a__	:					Any										=self.get_decoder()

        a__	:					Tuple										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

        a__	:					List[str]										="This is a test string"

        a__	:					Optional[int]										=processor(text=lowerCAmelCase__							)

        a__	:					Any										=tokenizer(lowerCAmelCase__							)

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







      def    _lowercase  (    self		,      lowerCAmelCase__=(2, 1_0, 1_6)		,      lowerCAmelCase__=7_7							)   ->						Dict:



        '''simple docstring'''




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







      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					Optional[int]										=self.get_feature_extractor()
        a__	:					Optional[Any]										=self.get_tokenizer()
        a__	:					Union[str, Any]										=self.get_decoder()

        a__	:					Tuple										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

        a__	:					int										=self._get_dummy_logits(shape=(1_0, 1_6)		,      seed=1_3							)

        a__	:					int										=processor.decode(lowerCAmelCase__							)

        a__	:					Optional[Any]										=decoder.decode_beams(lowerCAmelCase__							)[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    _lowercase  (    self		,      lowerCAmelCase__							)   ->						Dict:



        '''simple docstring'''




        a__	:					Dict										=self.get_feature_extractor()
        a__	:					Union[str, Any]										=self.get_tokenizer()
        a__	:					Union[str, Any]										=self.get_decoder()

        a__	:					Optional[int]										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

        a__	:					List[Any]										=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__	:					List[Any]										=processor.batch_decode(lowerCAmelCase__							)
        else:
          with get_context(lowerCAmelCase__							).Pool() as pool:
            a__	:					Tuple										=processor.batch_decode(lowerCAmelCase__		,      lowerCAmelCase__							)

        a__	:					Optional[int]										=list(lowerCAmelCase__							)

        with get_context("fork"							).Pool() as p:
          a__	:					Any										=decoder.decode_beams_batch(lowerCAmelCase__		,      lowerCAmelCase__							)

        a__					,							a__					,							a__	:					Any										=[], [], []
        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(lowerCAmelCase__		,      decoded_processor.text							)
        self.assertListEqual(["<s> <s> </s>", "<s> <s> <s>"]		,      decoded_processor.text							)
        self.assertListEqual(lowerCAmelCase__		,      decoded_processor.logit_score							)
        self.assertListEqual(lowerCAmelCase__		,      decoded_processor.lm_score							)







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



        '''simple docstring'''




        a__	:					Optional[Any]										=self.get_feature_extractor()
        a__	:					Optional[Any]										=self.get_tokenizer()
        a__	:					Union[str, Any]										=self.get_decoder()

        a__	:					List[str]										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

        a__	:					Optional[Any]										=self._get_dummy_logits()

        a__	:					List[Any]										=1_5
        a__	:					List[Any]										=-20.0
        a__	:					str										=-4.0

        a__	:					Union[str, Any]										=processor.batch_decode(
            lowerCAmelCase__		,      beam_width=lowerCAmelCase__		,      beam_prune_logp=lowerCAmelCase__		,      token_min_logp=lowerCAmelCase__		,      )
        a__	:					List[str]										=decoded_processor_out.text

        a__	:					List[Any]										=list(lowerCAmelCase__							)

        with get_context("fork"							).Pool() as pool:
          a__	:					Dict										=decoder.decode_beams_batch(
              lowerCAmelCase__		,      lowerCAmelCase__		,      beam_width=lowerCAmelCase__		,      beam_prune_logp=lowerCAmelCase__		,      token_min_logp=lowerCAmelCase__		,      )

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

        self.assertListEqual(lowerCAmelCase__		,      lowerCAmelCase__							)
        self.assertListEqual(["</s> <s> <s>", "<s> <s> <s>"]		,      lowerCAmelCase__							)

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

        self.assertTrue(np.array_equal(lowerCAmelCase__		,      decoded_processor_out.lm_score							)							)
        self.assertTrue(np.allclose([-15.5_54, -13.94_74]		,      lowerCAmelCase__		,      atol=1E-3							)							)







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Optional[int]										=self.get_feature_extractor()
        a__	:					List[Any]										=self.get_tokenizer()
        a__	:					Optional[int]										=self.get_decoder()

        a__	:					Dict										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

        a__	:					List[Any]										=self._get_dummy_logits()

        a__	:					Dict										=2.0
        a__	:					int										=5.0
        a__	:					str										=-20.0
        a__	:					Tuple										=True

        a__	:					Tuple										=processor.batch_decode(
            lowerCAmelCase__		,      alpha=lowerCAmelCase__		,      beta=lowerCAmelCase__		,      unk_score_offset=lowerCAmelCase__		,      lm_score_boundary=lowerCAmelCase__		,      )
        a__	:					Union[str, Any]										=decoded_processor_out.text

        a__	:					Tuple										=list(lowerCAmelCase__							)
        decoder.reset_params(
            alpha=lowerCAmelCase__		,      beta=lowerCAmelCase__		,      unk_score_offset=lowerCAmelCase__		,      lm_score_boundary=lowerCAmelCase__		,      )

        with get_context("fork"							).Pool() as pool:
          a__	:					str										=decoder.decode_beams_batch(
              lowerCAmelCase__		,      lowerCAmelCase__		,      )

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

        self.assertListEqual(lowerCAmelCase__		,      lowerCAmelCase__							)
        self.assertListEqual(["<s> </s> <s> </s> </s>", "</s> </s> <s> </s> </s>"]		,      lowerCAmelCase__							)
        a__	:					Dict										=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		,      lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					Any										=WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"							)

        a__	:					Dict										=processor.decoder.model_container[processor.decoder._model_key]
        a__	:					Dict										=Path(language_model._kenlm_model.path.decode("utf-8"							)							).parent.parent.absolute()

        a__	:					Union[str, Any]										=os.listdir(lowerCAmelCase__							)
        a__	:					Tuple										=["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(lowerCAmelCase__		,      lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					List[str]										=snapshot_download("hf-internal-testing/processor_with_lm"							)

        a__	:					Union[str, Any]										=WavaVecaProcessorWithLM.from_pretrained(lowerCAmelCase__							)

        a__	:					List[Any]										=processor.decoder.model_container[processor.decoder._model_key]
        a__	:					str										=Path(language_model._kenlm_model.path.decode("utf-8"							)							).parent.parent.absolute()

        a__	:					List[str]										=os.listdir(lowerCAmelCase__							)
        a__	:					Optional[Any]										=os.listdir(lowerCAmelCase__							)

        local_decoder_files.sort()
        expected_decoder_files.sort()

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







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



        '''simple docstring'''




        a__	:					Dict										=WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"							)
        a__	:					str										=AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"							)

        a__	:					str										=floats_list((3, 1_0_0_0)							)

        a__	:					int										=processor_wavaveca(lowerCAmelCase__		,      return_tensors="np"							)
        a__	:					List[Any]										=processor_auto(lowerCAmelCase__		,      return_tensors="np"							)

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

        a__	:					int										=self._get_dummy_logits()

        a__	:					Tuple										=processor_wavaveca.batch_decode(lowerCAmelCase__							)
        a__	:					int										=processor_auto.batch_decode(lowerCAmelCase__							)

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







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					List[str]										=self.get_feature_extractor()
        a__	:					List[Any]										=self.get_tokenizer()
        a__	:					Dict										=self.get_decoder()

        a__	:					Tuple										=WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase__		,      feature_extractor=lowerCAmelCase__		,      decoder=lowerCAmelCase__							)

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







      @staticmethod
      def    _lowercase  (    lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Any:



        '''simple docstring'''




        a__	:					List[Any]										=[d[key] for d in offsets]
        return retrieved_list







      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					Any										=WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"							)
        a__	:					Tuple										=self._get_dummy_logits()[0]

        a__	:					Union[str, Any]										=processor.decode(lowerCAmelCase__		,      output_word_offsets=lowerCAmelCase__							)
        # 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(lowerCAmelCase__		,      lowerCAmelCase__							)							)

        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    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Dict										=WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"							)
        a__	:					List[str]										=self._get_dummy_logits()

        a__	:					Union[str, Any]										=processor.batch_decode(lowerCAmelCase__		,      output_word_offsets=lowerCAmelCase__							)

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

        self.assertListEqual(
            [" ".join(self.get_from_offsets(lowerCAmelCase__		,      "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    _lowercase  (    self							)   ->						Tuple:



        '''simple docstring'''




        import torch

        a__	:					Optional[Any]										=load_dataset("common_voice"		,      "en"		,      split="train"		,      streaming=lowerCAmelCase__							)
        a__	:					int										=ds.cast_column("audio"		,      datasets.Audio(sampling_rate=1_6_0_0_0							)							)
        a__	:					List[str]										=iter(lowerCAmelCase__							)
        a__	:					Optional[int]										=next(lowerCAmelCase__							)

        a__	:					Optional[Any]										=AutoProcessor.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm"							)
        a__	:					Optional[int]										=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__	:					Any										=processor(sample["audio"]["array"]		,      return_tensors="pt"							).input_values

        with torch.no_grad():
          a__	:					Tuple										=model(lowerCAmelCase__							).logits.cpu().numpy()

        a__	:					Optional[int]										=processor.decode(logits[0]		,      output_word_offsets=lowerCAmelCase__							)

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

        a__	:					Tuple										="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(lowerCAmelCase__		,      "word"							)							)		,      lowerCAmelCase__							)
        self.assertEqual(" ".join(self.get_from_offsets(lowerCAmelCase__		,      "word"							)							)		,      output.text							)

        # output times
        a__	:					Dict										=torch.tensor(self.get_from_offsets(lowerCAmelCase__		,      "start_time"							)							)
        a__	:					str										=torch.tensor(self.get_from_offsets(lowerCAmelCase__		,      "end_time"							)							)

        # fmt: off
        a__	:					Dict										=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]							)
        a__	:					int										=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(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=0.01							)							)
        self.assertTrue(torch.allclose(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=0.01							)							)













import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


UpperCAmelCase	:	int		     =       {
    """cola""": 2,
    """mnli""": 3,
    """mrpc""": 2,
    """sst-2""": 2,
    """sts-b""": 1,
    """qqp""": 2,
    """qnli""": 2,
    """rte""": 2,
    """wnli""": 2,
}


logging.set_verbosity_info()
def      _A       (		SCREAMING_SNAKE_CASE     :   List[Any]       ,       SCREAMING_SNAKE_CASE     :   Dict       ,       SCREAMING_SNAKE_CASE     :   Dict       ,       SCREAMING_SNAKE_CASE     :   Union[str, Any]=None						):





  """simple docstring"""



  a__	:					Optional[int]										=XLNetConfig.from_json_file(SCREAMING_SNAKE_CASE						)

  a__	:					Dict										=finetuning_task.lower() if finetuning_task is not None else ""
  if finetuning_task in GLUE_TASKS_NUM_LABELS:
    print(f'''Building PyTorch XLNetForSequenceClassification model from configuration: {config}'''						)
    a__	:					List[str]										=finetuning_task
    a__	:					Tuple										=GLUE_TASKS_NUM_LABELS[finetuning_task]
    a__	:					List[Any]										=XLNetForSequenceClassification(SCREAMING_SNAKE_CASE						)
  elif "squad" in finetuning_task:
    a__	:					Optional[int]										=finetuning_task
    a__	:					Dict										=XLNetForQuestionAnswering(SCREAMING_SNAKE_CASE						)
  else:
    a__	:					List[Any]										=XLNetLMHeadModel(SCREAMING_SNAKE_CASE						)

  # Load weights from tf checkpoint
  load_tf_weights_in_xlnet(SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE						)

  # Save pytorch-model
  a__	:					Dict										=os.path.join(SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE						)
  a__	:					Dict										=os.path.join(SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE						)
  print(f'''Save PyTorch model to {os.path.abspath(SCREAMING_SNAKE_CASE						)}'''						)
  torch.save(model.state_dict()       ,       SCREAMING_SNAKE_CASE						)
  print(f'''Save configuration file to {os.path.abspath(SCREAMING_SNAKE_CASE						)}'''						)
  with open(SCREAMING_SNAKE_CASE       ,       "w"       ,       encoding="utf-8"						) as f:
    f.write(config.to_json_string()						)


if __name__ == "__main__":
     UpperCAmelCase	:	List[str]		     =       argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
     )
     parser.add_argument(
         """--xlnet_config_file""",
         default=None,
         type=str,
         required=True,
         help=(
             """The config json file corresponding to the pre-trained XLNet model. \n"""
             """This specifies the model architecture."""
         ),
     )
     parser.add_argument(
         """--pytorch_dump_folder_path""",
         default=None,
         type=str,
         required=True,
         help="""Path to the folder to store the PyTorch model or dataset/vocab.""",
     )
     parser.add_argument(
         """--finetuning_task""",
         default=None,
         type=str,
         help="""Name of a task on which the XLNet TensorFlow model was fine-tuned""",
     )
     UpperCAmelCase	:	int		     =       parser.parse_args()
     print(args)

     convert_xlnet_checkpoint_to_pytorch(
         args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
     )






