


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





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






         """simple docstring"""







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

         UpperCAmelCase							=     {
             """unk_token""": """<unk>""",
             """bos_token""": """<s>""",
             """eos_token""": """</s>""",
         }
         UpperCAmelCase							=     {
             """feature_size""": 1,
             """padding_value""": 0.0,
             """sampling_rate""": 1_6_0_0_0,
             """return_attention_mask""": False,
             """do_normalize""": True,
         }

         UpperCAmelCase							=     tempfile.mkdtemp()
         UpperCAmelCase							=     os.path.join(self.tmpdirname					,       VOCAB_FILES_NAMES["""vocab_file"""]       )
         UpperCAmelCase							=     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
         UpperCAmelCase							=     """hf-internal-testing/ngram-beam-search-decoder"""





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






         """simple docstring"""







         UpperCAmelCase							=     self.add_kwargs_tokens_map.copy()
         kwargs.update(__lowerCamelCase       )
         return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname					,       **__lowerCamelCase       )





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






         """simple docstring"""







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





      def        _lowercase				(		self						:					List[Any]					,       **__lowerCamelCase						:					int       )    ->   Tuple:






         """simple docstring"""







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





      def        _lowercase				(		self						:					Any       )    ->   str:






         """simple docstring"""







         shutil.rmtree(self.tmpdirname       )





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






         """simple docstring"""







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

         UpperCAmelCase							=     WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,       feature_extractor=__lowerCamelCase					,       decoder=__lowerCamelCase       )

         processor.save_pretrained(self.tmpdirname       )
         UpperCAmelCase							=     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						:					Any       )    ->   str:






         """simple docstring"""







         UpperCAmelCase							=     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
         UpperCAmelCase							=     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]       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     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						:					Tuple       )    ->   int:






         """simple docstring"""







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

         UpperCAmelCase							=     WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,       feature_extractor=__lowerCamelCase					,       decoder=__lowerCamelCase       )

         UpperCAmelCase							=     floats_list((3, 1_0_0_0)       )

         UpperCAmelCase							=     feature_extractor(__lowerCamelCase					,       return_tensors="""np"""       )
         UpperCAmelCase							=     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						:					Union[str, Any]       )    ->   List[Any]:






         """simple docstring"""







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

         UpperCAmelCase							=     WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,       feature_extractor=__lowerCamelCase					,       decoder=__lowerCamelCase       )

         UpperCAmelCase							=     """This is a test string"""

         UpperCAmelCase							=     processor(text=__lowerCamelCase       )

         UpperCAmelCase							=     tokenizer(__lowerCamelCase       )

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





      def        _lowercase				(		self						:					List[Any]					,       __lowerCamelCase						:					Optional[Any]=(2, 1_0, 1_6)					,       __lowerCamelCase						:					Union[str, Any]=7_7       )    ->   List[str]:






         """simple docstring"""







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





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






         """simple docstring"""







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

         UpperCAmelCase							=     WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,       feature_extractor=__lowerCamelCase					,       decoder=__lowerCamelCase       )

         UpperCAmelCase							=     self._get_dummy_logits(shape=(1_0, 1_6)					,       seed=1_3       )

         UpperCAmelCase							=     processor.decode(__lowerCamelCase       )

         UpperCAmelCase							=     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						:					List[Any]					,       __lowerCamelCase						:					int       )    ->   Any:






         """simple docstring"""







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

         UpperCAmelCase							=     WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,       feature_extractor=__lowerCamelCase					,       decoder=__lowerCamelCase       )

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

         UpperCAmelCase							=     list(__lowerCamelCase       )

         with get_context("""fork"""       ).Pool() as p:
            UpperCAmelCase							=     decoder.decode_beams_batch(__lowerCamelCase					,       __lowerCamelCase       )

         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     [], [], []
         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						:					int       )    ->   int:






         """simple docstring"""







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

         UpperCAmelCase							=     WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,       feature_extractor=__lowerCamelCase					,       decoder=__lowerCamelCase       )

         UpperCAmelCase							=     self._get_dummy_logits()

         UpperCAmelCase							=     1_5
         UpperCAmelCase							=     -20.0
         UpperCAmelCase							=     -4.0

         UpperCAmelCase							=     processor.batch_decode(
             __lowerCamelCase					,       beam_width=__lowerCamelCase					,       beam_prune_logp=__lowerCamelCase					,       token_min_logp=__lowerCamelCase					,       )
         UpperCAmelCase							=     decoded_processor_out.text

         UpperCAmelCase							=     list(__lowerCamelCase       )

         with get_context("""fork"""       ).Pool() as pool:
            UpperCAmelCase							=     decoder.decode_beams_batch(
                __lowerCamelCase					,       __lowerCamelCase					,       beam_width=__lowerCamelCase					,       beam_prune_logp=__lowerCamelCase					,       token_min_logp=__lowerCamelCase					,       )

         UpperCAmelCase							=     [d[0][0] for d in decoded_decoder_out]
         UpperCAmelCase							=     [d[0][2] for d in decoded_decoder_out]
         UpperCAmelCase							=     [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.054, -18.447]					,       __lowerCamelCase					,       atol=1e-3       )       )

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





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






         """simple docstring"""







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

         UpperCAmelCase							=     WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,       feature_extractor=__lowerCamelCase					,       decoder=__lowerCamelCase       )

         UpperCAmelCase							=     self._get_dummy_logits()

         UpperCAmelCase							=     2.0
         UpperCAmelCase							=     5.0
         UpperCAmelCase							=     -20.0
         UpperCAmelCase							=     True

         UpperCAmelCase							=     processor.batch_decode(
             __lowerCamelCase					,       alpha=__lowerCamelCase					,       beta=__lowerCamelCase					,       unk_score_offset=__lowerCamelCase					,       lm_score_boundary=__lowerCamelCase					,       )
         UpperCAmelCase							=     decoded_processor_out.text

         UpperCAmelCase							=     list(__lowerCamelCase       )
         decoder.reset_params(
             alpha=__lowerCamelCase					,       beta=__lowerCamelCase					,       unk_score_offset=__lowerCamelCase					,       lm_score_boundary=__lowerCamelCase					,       )

         with get_context("""fork"""       ).Pool() as pool:
            UpperCAmelCase							=     decoder.decode_beams_batch(
                __lowerCamelCase					,       __lowerCamelCase					,       )

         UpperCAmelCase							=     [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       )
         UpperCAmelCase							=     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						:					Optional[int]       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase							=     WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""       )

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

         UpperCAmelCase							=     os.listdir(__lowerCamelCase       )
         UpperCAmelCase							=     ["""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]       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase							=     snapshot_download("""hf-internal-testing/processor_with_lm"""       )

         UpperCAmelCase							=     WavaVecaProcessorWithLM.from_pretrained(__lowerCamelCase       )

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

         UpperCAmelCase							=     os.listdir(__lowerCamelCase       )
         UpperCAmelCase							=     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						:					Dict       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""       )
         UpperCAmelCase							=     AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""       )

         UpperCAmelCase							=     floats_list((3, 1_0_0_0)       )

         UpperCAmelCase							=     processor_wavaveca(__lowerCamelCase					,       return_tensors="""np"""       )
         UpperCAmelCase							=     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       )

         UpperCAmelCase							=     self._get_dummy_logits()

         UpperCAmelCase							=     processor_wavaveca.batch_decode(__lowerCamelCase       )
         UpperCAmelCase							=     processor_auto.batch_decode(__lowerCamelCase       )

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





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






         """simple docstring"""







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

         UpperCAmelCase							=     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						:					str					,       __lowerCamelCase						:					Tuple       )    ->   Optional[Any]:






         """simple docstring"""







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





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






         """simple docstring"""







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

         UpperCAmelCase							=     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						:					Union[str, Any]       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""       )
         UpperCAmelCase							=     self._get_dummy_logits()

         UpperCAmelCase							=     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						:					Dict       )    ->   List[Any]:






         """simple docstring"""







         import torch

         UpperCAmelCase							=     load_dataset("""common_voice"""					,       """en"""					,       split="""train"""					,       streaming=__lowerCamelCase       )
         UpperCAmelCase							=     ds.cast_column("""audio"""					,       datasets.Audio(sampling_rate=1_6_0_0_0       )       )
         UpperCAmelCase							=     iter(__lowerCamelCase       )
         UpperCAmelCase							=     next(__lowerCamelCase       )

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

         with torch.no_grad():
            UpperCAmelCase							=     model(__lowerCamelCase       ).logits.cpu().numpy()

         UpperCAmelCase							=     processor.decode(logits[0]					,       output_word_offsets=__lowerCamelCase       )

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

         UpperCAmelCase							=     """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
         UpperCAmelCase							=     torch.tensor(self.get_from_offsets(__lowerCamelCase					,       """start_time"""       )       )
         UpperCAmelCase							=     torch.tensor(self.get_from_offsets(__lowerCamelCase					,       """end_time"""       )       )

         # fmt: off
         UpperCAmelCase							=     torch.tensor([1.4_199, 1.6_599, 2.2_599, 3.0, 3.24, 3.5_999, 3.7_999, 4.0_999, 4.26, 4.94, 5.28, 5.6_599, 5.78, 5.94, 6.32, 6.5_399, 6.6_599]       )
         UpperCAmelCase							=     torch.tensor([1.5_399, 1.8_999, 2.9, 3.16, 3.5_399, 3.72, 4.0_199, 4.1_799, 4.76, 5.1_599, 5.5_599, 5.6_999, 5.86, 6.1_999, 6.38, 6.6_199, 6.94]       )
         # fmt: on

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







from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import (
     TFDeiTForImageClassification,
     TFDeiTForImageClassificationWithTeacher,
     TFDeiTForMaskedImageModeling,
     TFDeiTModel,
 )
 from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import DeiTImageProcessor






class   __lowercase				:





      def __init__(		self						:					str					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Optional[Any]=1_3					,       __lowerCamelCase						:					Optional[Any]=3_0					,       __lowerCamelCase						:					Any=2					,       __lowerCamelCase						:					List[str]=3					,       __lowerCamelCase						:					str=True					,       __lowerCamelCase						:					str=True					,       __lowerCamelCase						:					Union[str, Any]=3_2					,       __lowerCamelCase						:					Optional[int]=2					,       __lowerCamelCase						:					Optional[int]=4					,       __lowerCamelCase						:					Union[str, Any]=3_7					,       __lowerCamelCase						:					str="gelu"					,       __lowerCamelCase						:					str=0.1					,       __lowerCamelCase						:					List[str]=0.1					,       __lowerCamelCase						:					str=1_0					,       __lowerCamelCase						:					Union[str, Any]=0.02					,       __lowerCamelCase						:					List[str]=3					,       __lowerCamelCase						:					Optional[Any]=None					,       __lowerCamelCase						:					Any=2					,       )    ->   Tuple:






         """simple docstring"""







         UpperCAmelCase							=     parent
         UpperCAmelCase							=     batch_size
         UpperCAmelCase							=     image_size
         UpperCAmelCase							=     patch_size
         UpperCAmelCase							=     num_channels
         UpperCAmelCase							=     is_training
         UpperCAmelCase							=     use_labels
         UpperCAmelCase							=     hidden_size
         UpperCAmelCase							=     num_hidden_layers
         UpperCAmelCase							=     num_attention_heads
         UpperCAmelCase							=     intermediate_size
         UpperCAmelCase							=     hidden_act
         UpperCAmelCase							=     hidden_dropout_prob
         UpperCAmelCase							=     attention_probs_dropout_prob
         UpperCAmelCase							=     type_sequence_label_size
         UpperCAmelCase							=     initializer_range
         UpperCAmelCase							=     scope
         UpperCAmelCase							=     encoder_stride

         # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
         UpperCAmelCase							=     (image_size // patch_size) ** 2
         UpperCAmelCase							=     num_patches + 2





      def        _lowercase				(		self						:					Tuple       )    ->   int:






         """simple docstring"""







         UpperCAmelCase							=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

         UpperCAmelCase							=     None
         if self.use_labels:
            UpperCAmelCase							=     ids_tensor([self.batch_size]					,       self.type_sequence_label_size       )

         UpperCAmelCase							=     self.get_config()

         return config, pixel_values, labels





      def        _lowercase				(		self						:					Tuple       )    ->   str:






         """simple docstring"""







         return DeiTConfig(
             image_size=self.image_size					,       patch_size=self.patch_size					,       num_channels=self.num_channels					,       hidden_size=self.hidden_size					,       num_hidden_layers=self.num_hidden_layers					,       num_attention_heads=self.num_attention_heads					,       intermediate_size=self.intermediate_size					,       hidden_act=self.hidden_act					,       hidden_dropout_prob=self.hidden_dropout_prob					,       attention_probs_dropout_prob=self.attention_probs_dropout_prob					,       is_decoder=__lowerCamelCase					,       initializer_range=self.initializer_range					,       encoder_stride=self.encoder_stride					,       )





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






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTModel(config=__lowerCamelCase       )
         UpperCAmelCase							=     model(__lowerCamelCase       )
         self.parent.assertEqual(result.last_hidden_state.shape					,       (self.batch_size, self.seq_length, self.hidden_size)       )





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






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTForMaskedImageModeling(config=__lowerCamelCase       )
         UpperCAmelCase							=     model(__lowerCamelCase       )
         self.parent.assertEqual(
             result.reconstruction.shape					,       (self.batch_size, self.num_channels, self.image_size, self.image_size)       )

         # test greyscale images
         UpperCAmelCase							=     1
         UpperCAmelCase							=     TFDeiTForMaskedImageModeling(__lowerCamelCase       )

         UpperCAmelCase							=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
         UpperCAmelCase							=     model(__lowerCamelCase       )
         self.parent.assertEqual(result.reconstruction.shape					,       (self.batch_size, 1, self.image_size, self.image_size)       )





      def        _lowercase				(		self						:					Union[str, Any]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					Dict					,       __lowerCamelCase						:					int       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     self.type_sequence_label_size
         UpperCAmelCase							=     TFDeiTForImageClassification(__lowerCamelCase       )
         UpperCAmelCase							=     model(__lowerCamelCase					,       labels=__lowerCamelCase       )
         self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.type_sequence_label_size)       )

         # test greyscale images
         UpperCAmelCase							=     1
         UpperCAmelCase							=     TFDeiTForImageClassification(__lowerCamelCase       )

         UpperCAmelCase							=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
         UpperCAmelCase							=     model(__lowerCamelCase					,       labels=__lowerCamelCase       )
         self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.type_sequence_label_size)       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.prepare_config_and_inputs()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     config_and_inputs
         UpperCAmelCase							=     {"""pixel_values""": pixel_values}
         return config, inputs_dict






@require_tf
class   __lowercase				(  __snake_case ,					__snake_case ,					unittest.TestCase  ):
      UpperCamelCase								=       (
          (
              TFDeiTModel,
              TFDeiTForImageClassification,
              TFDeiTForImageClassificationWithTeacher,
              TFDeiTForMaskedImageModeling,
          )
          if is_tf_available()
          else ()
      )
      UpperCamelCase								=       (
          {
              '''feature-extraction''': TFDeiTModel,
              '''image-classification''': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
          }
          if is_tf_available()
          else {}
      )

      UpperCamelCase								=       False
      UpperCamelCase								=       False
      UpperCamelCase								=       False
      UpperCamelCase								=       False





      def        _lowercase				(		self						:					str       )    ->   str:






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTModelTester(self       )
         UpperCAmelCase							=     ConfigTester(self					,       config_class=__lowerCamelCase					,       has_text_modality=__lowerCamelCase					,       hidden_size=3_7       )





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






         """simple docstring"""







         self.config_tester.run_common_tests()





      @unittest.skip(reason="""DeiT does not use inputs_embeds"""       )
      def        _lowercase				(		self						:					List[Any]       )    ->   Dict:






         """simple docstring"""







         pass





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






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            UpperCAmelCase							=     model_class(__lowerCamelCase       )
            self.assertIsInstance(model.get_input_embeddings()					,       (tf.keras.layers.Layer)       )
            UpperCAmelCase							=     model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(__lowerCamelCase					,       tf.keras.layers.Dense       )       )





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






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            UpperCAmelCase							=     model_class(__lowerCamelCase       )
            UpperCAmelCase							=     inspect.signature(model.call       )
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            UpperCAmelCase							=     [*signature.parameters.keys()]

            UpperCAmelCase							=     ["""pixel_values"""]
            self.assertListEqual(arg_names[:1]					,       __lowerCamelCase       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCamelCase       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*__lowerCamelCase       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase       )





      def        _lowercase				(		self						:					Optional[int]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					Any=False       )    ->   int:






         """simple docstring"""







         UpperCAmelCase							=     super()._prepare_for_class(__lowerCamelCase					,       __lowerCamelCase					,       return_labels=__lowerCamelCase       )

         if return_labels:
            if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call       ).parameters:
               del inputs_dict["labels"]

         return inputs_dict





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






         """simple docstring"""







         for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCAmelCase							=     TFDeiTModel.from_pretrained(__lowerCamelCase       )
            self.assertIsNotNone(__lowerCamelCase       )

def   _UpperCamelCase				(						)						->Tuple:
   UpperCAmelCase							=     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
   return image







@require_tf
@require_vision
class   __lowercase				(  unittest.TestCase  ):





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






         """simple docstring"""







         return (
             DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""       )
             if is_vision_available()
             else None
         )





      @slow
      def        _lowercase				(		self						:					Optional[int]       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""       )

         UpperCAmelCase							=     self.default_image_processor
         UpperCAmelCase							=     prepare_img()
         UpperCAmelCase							=     image_processor(images=__lowerCamelCase					,       return_tensors="""tf"""       )

         # forward pass
         UpperCAmelCase							=     model(**__lowerCamelCase       )

         # verify the logits
         UpperCAmelCase							=     tf.TensorShape((1, 1_0_0_0)       )
         self.assertEqual(outputs.logits.shape					,       __lowerCamelCase       )

         UpperCAmelCase							=     tf.constant([-1.0_266, 0.1_912, -1.2_861]       )

         self.assertTrue(np.allclose(outputs.logits[0, :3]					,       __lowerCamelCase					,       atol=1e-4       )       )



