




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 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



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

   _snake_case							:	Optional[Any]			=						{
       '''unk_token''': '''<unk>''',
       '''bos_token''': '''<s>''',
       '''eos_token''': '''</s>''',
   }
   _snake_case							:	List[Any]			=						{
       '''feature_size''': 1,
       '''padding_value''': 0.0,
       '''sampling_rate''': 16_000,
       '''return_attention_mask''': False,
       '''do_normalize''': True,
   }

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

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

   # load decoder from hub
   _snake_case							:	Dict			=						'''hf-internal-testing/ngram-beam-search-decoder'''


 def 				UpperCamelCase_     (       self:	Optional[int],   **a_:	Any	):




   '''simple docstring'''



   _snake_case							:	Dict			=						self.add_kwargs_tokens_map.copy()
   kwargs.update(__snake_case	)
   return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname,   **__snake_case	)


 def 				UpperCamelCase_     (       self:	Optional[int],   **a_:	List[Any]	):




   '''simple docstring'''



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


 def 				UpperCamelCase_     (       self:	Union[str, Any],   **a_:	List[Any]	):




   '''simple docstring'''



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


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   shutil.rmtree(self.tmpdirname	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	Dict			=						self.get_tokenizer()
   _snake_case							:	Tuple			=						self.get_feature_extractor()
   _snake_case							:	Optional[int]			=						self.get_decoder()

   _snake_case							:	str			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

   processor.save_pretrained(self.tmpdirname	)
   _snake_case							:	Union[str, Any]			=						WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname	)

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

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

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


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   _snake_case							:	Dict			=						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
   _snake_case							:	Any			=						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 				UpperCamelCase_     (       self:	Optional[int]	):




   '''simple docstring'''



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


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



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

   _snake_case							:	Any			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

   _snake_case							:	Tuple			=						floats_list((3, 1_000)	)

   _snake_case							:	Union[str, Any]			=						feature_extractor(__snake_case,   return_tensors="""np"""	)
   _snake_case							:	Optional[Any]			=						processor(__snake_case,   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 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



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

   _snake_case							:	List[str]			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

   _snake_case							:	Union[str, Any]			=						'''This is a test string'''

   _snake_case							:	Optional[int]			=						processor(text=__snake_case	)

   _snake_case							:	Dict			=						tokenizer(__snake_case	)

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


 def 				UpperCamelCase_     (       self:	Dict,   a_:	List[str]=(2, 10, 16),   a_:	Tuple=77	):




   '''simple docstring'''



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


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	Any			=						self.get_feature_extractor()
   _snake_case							:	int			=						self.get_tokenizer()
   _snake_case							:	int			=						self.get_decoder()

   _snake_case							:	int			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

   _snake_case							:	Dict			=						self._get_dummy_logits(shape=(10, 16),   seed=13	)

   _snake_case							:	Optional[Any]			=						processor.decode(__snake_case	)

   _snake_case							:	Optional[int]			=						decoder.decode_beams(__snake_case	)[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 				UpperCamelCase_     (       self:	Any,   a_:	str	):




   '''simple docstring'''



   _snake_case							:	str			=						self.get_feature_extractor()
   _snake_case							:	str			=						self.get_tokenizer()
   _snake_case							:	Any			=						self.get_decoder()

   _snake_case							:	Any			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

   _snake_case							:	Tuple			=						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:
     _snake_case							:	int			=						processor.batch_decode(__snake_case	)
   else:
     with get_context(__snake_case	).Pool() as pool:
       _snake_case							:	List[str]			=						processor.batch_decode(__snake_case,   __snake_case	)

   _snake_case							:	List[Any]			=						list(__snake_case	)

   with get_context("""fork"""	).Pool() as p:
     _snake_case							:	int			=						decoder.decode_beams_batch(__snake_case,   __snake_case	)

   _snake_case							:	Optional[int]			=						[], [], []
   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(__snake_case,   decoded_processor.text	)
   self.assertListEqual(["""<s> <s> </s>""", """<s> <s> <s>"""],   decoded_processor.text	)
   self.assertListEqual(__snake_case,   decoded_processor.logit_score	)
   self.assertListEqual(__snake_case,   decoded_processor.lm_score	)


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



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

   _snake_case							:	str			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

   _snake_case							:	Dict			=						self._get_dummy_logits()

   _snake_case							:	Optional[int]			=						15
   _snake_case							:	Optional[Any]			=						-20.0
   _snake_case							:	int			=						-4.0

   _snake_case							:	int			=						processor.batch_decode(
       __snake_case,   beam_width=__snake_case,   beam_prune_logp=__snake_case,   token_min_logp=__snake_case,   )
   _snake_case							:	List[str]			=						decoded_processor_out.text

   _snake_case							:	List[str]			=						list(__snake_case	)

   with get_context("""fork"""	).Pool() as pool:
     _snake_case							:	List[Any]			=						decoder.decode_beams_batch(
         __snake_case,   __snake_case,   beam_width=__snake_case,   beam_prune_logp=__snake_case,   token_min_logp=__snake_case,   )

   _snake_case							:	Dict			=						[d[0][0] for d in decoded_decoder_out]
   _snake_case							:	Optional[int]			=						[d[0][2] for d in decoded_decoder_out]
   _snake_case							:	Tuple			=						[d[0][3] for d in decoded_decoder_out]

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

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

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


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



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

   _snake_case							:	str			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

   _snake_case							:	Dict			=						self._get_dummy_logits()

   _snake_case							:	str			=						2.0
   _snake_case							:	Any			=						5.0
   _snake_case							:	Optional[Any]			=						-20.0
   _snake_case							:	List[Any]			=						True

   _snake_case							:	List[Any]			=						processor.batch_decode(
       __snake_case,   alpha=__snake_case,   beta=__snake_case,   unk_score_offset=__snake_case,   lm_score_boundary=__snake_case,   )
   _snake_case							:	Dict			=						decoded_processor_out.text

   _snake_case							:	List[Any]			=						list(__snake_case	)
   decoder.reset_params(
       alpha=__snake_case,   beta=__snake_case,   unk_score_offset=__snake_case,   lm_score_boundary=__snake_case,   )

   with get_context("""fork"""	).Pool() as pool:
     _snake_case							:	Dict			=						decoder.decode_beams_batch(
         __snake_case,   __snake_case,   )

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

   self.assertListEqual(__snake_case,   __snake_case	)
   self.assertListEqual(["""<s> </s> <s> </s> </s>""", """</s> </s> <s> </s> </s>"""],   __snake_case	)
   _snake_case							:	List[Any]			=						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,   __snake_case	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	Dict			=						WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""	)

   _snake_case							:	Optional[Any]			=						processor.decoder.model_container[processor.decoder._model_key]
   _snake_case							:	Tuple			=						Path(language_model._kenlm_model.path.decode("""utf-8"""	)	).parent.parent.absolute()

   _snake_case							:	Any			=						os.listdir(__snake_case	)
   _snake_case							:	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(__snake_case,   __snake_case	)


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						snapshot_download("""hf-internal-testing/processor_with_lm"""	)

   _snake_case							:	Tuple			=						WavaVecaProcessorWithLM.from_pretrained(__snake_case	)

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

   _snake_case							:	List[Any]			=						os.listdir(__snake_case	)
   _snake_case							:	Any			=						os.listdir(__snake_case	)

   local_decoder_files.sort()
   expected_decoder_files.sort()

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


 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""	)
   _snake_case							:	List[Any]			=						AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""	)

   _snake_case							:	List[Any]			=						floats_list((3, 1_000)	)

   _snake_case							:	Optional[Any]			=						processor_wavaveca(__snake_case,   return_tensors="""np"""	)
   _snake_case							:	int			=						processor_auto(__snake_case,   return_tensors="""np"""	)

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

   _snake_case							:	Union[str, Any]			=						self._get_dummy_logits()

   _snake_case							:	Optional[Any]			=						processor_wavaveca.batch_decode(__snake_case	)
   _snake_case							:	Optional[Any]			=						processor_auto.batch_decode(__snake_case	)

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


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						self.get_feature_extractor()
   _snake_case							:	int			=						self.get_tokenizer()
   _snake_case							:	Dict			=						self.get_decoder()

   _snake_case							:	str			=						WavaVecaProcessorWithLM(tokenizer=__snake_case,   feature_extractor=__snake_case,   decoder=__snake_case	)

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


 @staticmethod
 def 				UpperCamelCase_     (       a_:	Optional[Any],   a_:	Union[str, Any]	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						[d[key] for d in offsets]
   return retrieved_list


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



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

   _snake_case							:	Any			=						processor.decode(__snake_case,   output_word_offsets=__snake_case	)
   # 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(__snake_case,   __snake_case	)	)

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




   '''simple docstring'''



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

   _snake_case							:	List[str]			=						processor.batch_decode(__snake_case,   output_word_offsets=__snake_case	)

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

   self.assertListEqual(
       [""" """.join(self.get_from_offsets(__snake_case,   """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 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   import torch

   _snake_case							:	Union[str, Any]			=						load_dataset("""common_voice""",   """en""",   split="""train""",   streaming=__snake_case	)
   _snake_case							:	Optional[Any]			=						ds.cast_column("""audio""",   datasets.Audio(sampling_rate=16_000	)	)
   _snake_case							:	Dict			=						iter(__snake_case	)
   _snake_case							:	Tuple			=						next(__snake_case	)

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

   with torch.no_grad():
     _snake_case							:	Dict			=						model(__snake_case	).logits.cpu().numpy()

   _snake_case							:	Dict			=						processor.decode(logits[0],   output_word_offsets=__snake_case	)

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

   _snake_case							:	Dict			=						'''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(__snake_case,   """word"""	)	),   __snake_case	)
   self.assertEqual(""" """.join(self.get_from_offsets(__snake_case,   """word"""	)	),   output.text	)

   # output times
   _snake_case							:	List[Any]			=						torch.tensor(self.get_from_offsets(__snake_case,   """start_time"""	)	)
   _snake_case							:	Any			=						torch.tensor(self.get_from_offsets(__snake_case,   """end_time"""	)	)

   # fmt: off
   _snake_case							:	Union[str, Any]			=						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]	)
   _snake_case							:	Optional[Any]			=						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(__snake_case,   __snake_case,   atol=0.01	)	)
   self.assertTrue(torch.allclose(__snake_case,   __snake_case,   atol=0.01	)	)








"""simple docstring"""


import argparse
import collections
import json
from pathlib import Path

import requests
import torch
import yaml
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTImageProcessor,
    MobileViTVaConfig,
    MobileViTVaForImageClassification,
    MobileViTVaForSemanticSegmentation,
)
from transformers.utils import logging


logging.set_verbosity_info()
A_      			=       logging.get_logger(__name__)


def       UpperCAmelCase__ (snake_case__							:      Optional[int]					):
  """simple docstring"""

  print("""Loading config file..."""					)

  def flatten_yaml_as_dict(snake_case__							:      List[Any]	,     snake_case__							:      Optional[Any]=""	,     snake_case__							:      Tuple="."					):
    _snake_case							:	Union[str, Any]			=						[]
    for k, v in d.items():
      _snake_case							:	List[str]			=						parent_key + sep + k if parent_key else k
      if isinstance(snake_case__	,     collections.abc.MutableMapping					):
        items.extend(flatten_yaml_as_dict(snake_case__	,     snake_case__	,     sep=snake_case__					).items()					)
      else:
        items.append((new_key, v)					)
    return dict(snake_case__					)

  _snake_case							:	Dict			=						argparse.Namespace()
  with open(snake_case__	,     """r"""					) as yaml_file:
    try:
      _snake_case							:	List[Any]			=						yaml.load(snake_case__	,     Loader=yaml.FullLoader					)

      _snake_case							:	Any			=						flatten_yaml_as_dict(snake_case__					)
      for k, v in flat_cfg.items():
        setattr(snake_case__	,     snake_case__	,     snake_case__					)
    except yaml.YAMLError as exc:
      logger.error("""Error while loading config file: {}. Error message: {}""".format(snake_case__	,     str(snake_case__					)					)					)
  return config


def       UpperCAmelCase__ (snake_case__							:      int	,     snake_case__							:      int					):
  """simple docstring"""

  _snake_case							:	Dict			=						MobileViTVaConfig()

  _snake_case							:	Optional[int]			=						False

  # dataset
  if task_name.startswith("""imagenet1k_"""					):
    _snake_case							:	Dict			=						10_00
    if int(task_name.strip().split("""_"""					)[-1]					) == 3_84:
      _snake_case							:	Union[str, Any]			=						3_84
    else:
      _snake_case							:	Optional[Any]			=						2_56
    _snake_case							:	str			=						"""imagenet-1k-id2label.json"""
  elif task_name.startswith("""imagenet21k_to_1k_"""					):
    _snake_case							:	str			=						2_10_00
    if int(task_name.strip().split("""_"""					)[-1]					) == 3_84:
      _snake_case							:	Dict			=						3_84
    else:
      _snake_case							:	Union[str, Any]			=						2_56
    _snake_case							:	Tuple			=						"""imagenet-22k-id2label.json"""
  elif task_name.startswith("""ade20k_"""					):
    _snake_case							:	Tuple			=						1_51
    _snake_case							:	str			=						5_12
    _snake_case							:	List[Any]			=						"""ade20k-id2label.json"""
    _snake_case							:	Union[str, Any]			=						True
  elif task_name.startswith("""voc_"""					):
    _snake_case							:	List[Any]			=						21
    _snake_case							:	List[str]			=						5_12
    _snake_case							:	int			=						"""pascal-voc-id2label.json"""
    _snake_case							:	int			=						True

  # orig_config
  _snake_case							:	int			=						load_orig_config_file(snake_case__					)
  assert getattr(snake_case__	,     """model.classification.name"""	,     -1					) == "mobilevit_v2", "Invalid model"
  _snake_case							:	str			=						getattr(snake_case__	,     """model.classification.mitv2.width_multiplier"""	,     1.0					)
  assert (
      getattr(snake_case__	,     """model.classification.mitv2.attn_norm_layer"""	,     -1					) == "layer_norm_2d"
  ), "Norm layers other than layer_norm_2d is not supported"
  _snake_case							:	int			=						getattr(snake_case__	,     """model.classification.activation.name"""	,     """swish"""					)
  # config.image_size == getattr(orig_config,  'sampler.bs.crop_size_width', 256)

  if is_segmentation_model:
    _snake_case							:	Tuple			=						getattr(snake_case__	,     """model.segmentation.output_stride"""	,     16					)
    if "_deeplabv3" in task_name:
      _snake_case							:	Any			=						getattr(snake_case__	,     """model.segmentation.deeplabv3.aspp_rates"""	,     [12, 24, 36]					)
      _snake_case							:	Tuple			=						getattr(snake_case__	,     """model.segmentation.deeplabv3.aspp_out_channels"""	,     5_12					)
      _snake_case							:	Any			=						getattr(snake_case__	,     """model.segmentation.deeplabv3.aspp_dropout"""	,     0.1					)

    # id2label
  _snake_case							:	Union[str, Any]			=						"""huggingface/label-files"""
  _snake_case							:	Any			=						json.load(open(hf_hub_download(snake_case__	,     snake_case__	,     repo_type="""dataset"""					)	,     """r"""					)					)
  _snake_case							:	List[Any]			=						{int(snake_case__					): v for k, v in idalabel.items()}
  _snake_case							:	Tuple			=						idalabel
  _snake_case							:	Any			=						{v: k for k, v in idalabel.items()}

  return config


def       UpperCAmelCase__ (snake_case__							:      Optional[Any]	,     snake_case__							:      Tuple	,     snake_case__							:      List[Any]					):
  """simple docstring"""

  _snake_case							:	List[str]			=						dct.pop(snake_case__					)
  _snake_case							:	List[Any]			=						val


def       UpperCAmelCase__ (snake_case__							:      Optional[Any]	,     snake_case__							:      int=False					):
  """simple docstring"""

  if base_model:
    _snake_case							:	Any			=						""""""
  else:
    _snake_case							:	Union[str, Any]			=						"""mobilevitv2."""

  _snake_case							:	Dict			=						[]
  for k in state_dict.keys():
    if k[:8] == "encoder.":
      _snake_case							:	List[str]			=						k[8:]
    else:
      _snake_case							:	str			=						k

    if ".block." in k:
      _snake_case							:	Optional[int]			=						k_new.replace(""".block."""	,     """."""					)
    if ".conv." in k:
      _snake_case							:	Union[str, Any]			=						k_new.replace(""".conv."""	,     """.convolution."""					)
    if ".norm." in k:
      _snake_case							:	str			=						k_new.replace(""".norm."""	,     """.normalization."""					)

    if "conv_1." in k:
      _snake_case							:	int			=						k_new.replace("""conv_1."""	,     F"{model_prefix}conv_stem."					)
    for i in [1, 2]:
      if F"layer_{i}." in k:
        _snake_case							:	Tuple			=						k_new.replace(F"layer_{i}."	,     F"{model_prefix}encoder.layer.{i-1}.layer."					)
    if ".exp_1x1." in k:
      _snake_case							:	Optional[Any]			=						k_new.replace(""".exp_1x1."""	,     """.expand_1x1."""					)
    if ".red_1x1." in k:
      _snake_case							:	Optional[Any]			=						k_new.replace(""".red_1x1."""	,     """.reduce_1x1."""					)

    for i in [3, 4, 5]:
      if F"layer_{i}.0." in k:
        _snake_case							:	Tuple			=						k_new.replace(F"layer_{i}.0."	,     F"{model_prefix}encoder.layer.{i-1}.downsampling_layer."					)
      if F"layer_{i}.1.local_rep.0." in k:
        _snake_case							:	Any			=						k_new.replace(F"layer_{i}.1.local_rep.0."	,     F"{model_prefix}encoder.layer.{i-1}.conv_kxk."					)
      if F"layer_{i}.1.local_rep.1." in k:
        _snake_case							:	str			=						k_new.replace(F"layer_{i}.1.local_rep.1."	,     F"{model_prefix}encoder.layer.{i-1}.conv_1x1."					)

    for i in [3, 4, 5]:
      if i == 3:
        _snake_case							:	Optional[Any]			=						[0, 1]
      elif i == 4:
        _snake_case							:	Any			=						[0, 1, 2, 3]
      elif i == 5:
        _snake_case							:	List[Any]			=						[0, 1, 2]

      for j in j_in:
        if F"layer_{i}.1.global_rep.{j}." in k:
          _snake_case							:	Any			=						k_new.replace(
              F"layer_{i}.1.global_rep.{j}."	,     F"{model_prefix}encoder.layer.{i-1}.transformer.layer.{j}."					)
      if F"layer_{i}.1.global_rep.{j+1}." in k:
        _snake_case							:	List[Any]			=						k_new.replace(
            F"layer_{i}.1.global_rep.{j+1}."	,     F"{model_prefix}encoder.layer.{i-1}.layernorm."					)

      if F"layer_{i}.1.conv_proj." in k:
        _snake_case							:	Union[str, Any]			=						k_new.replace(F"layer_{i}.1.conv_proj."	,     F"{model_prefix}encoder.layer.{i-1}.conv_projection."					)

    if "pre_norm_attn.0." in k:
      _snake_case							:	List[Any]			=						k_new.replace("""pre_norm_attn.0."""	,     """layernorm_before."""					)
    if "pre_norm_attn.1." in k:
      _snake_case							:	Optional[int]			=						k_new.replace("""pre_norm_attn.1."""	,     """attention."""					)
    if "pre_norm_ffn.0." in k:
      _snake_case							:	List[Any]			=						k_new.replace("""pre_norm_ffn.0."""	,     """layernorm_after."""					)
    if "pre_norm_ffn.1." in k:
      _snake_case							:	Tuple			=						k_new.replace("""pre_norm_ffn.1."""	,     """ffn.conv1."""					)
    if "pre_norm_ffn.3." in k:
      _snake_case							:	Any			=						k_new.replace("""pre_norm_ffn.3."""	,     """ffn.conv2."""					)

    if "classifier.1." in k:
      _snake_case							:	List[str]			=						k_new.replace("""classifier.1."""	,     """classifier."""					)

    if "seg_head." in k:
      _snake_case							:	str			=						k_new.replace("""seg_head."""	,     """segmentation_head."""					)
    if ".aspp_layer." in k:
      _snake_case							:	Optional[int]			=						k_new.replace(""".aspp_layer."""	,     """."""					)
    if ".aspp_pool." in k:
      _snake_case							:	int			=						k_new.replace(""".aspp_pool."""	,     """."""					)

    rename_keys.append((k, k_new)					)
  return rename_keys


def       UpperCAmelCase__ (snake_case__							:      str					):
  """simple docstring"""

  _snake_case							:	List[str]			=						[]
  for k in state_dict.keys():
    if k.startswith("""seg_head.aux_head."""					):
      keys_to_ignore.append(snake_case__					)
  for k in keys_to_ignore:
    state_dict.pop(snake_case__	,     snake_case__					)


def       UpperCAmelCase__ ():
  """simple docstring"""

  _snake_case							:	str			=						"""http://images.cocodataset.org/val2017/000000039769.jpg"""
  # url = "https://cdn.britannica.com/86/141086-050-9D7C75EE/Gulfstream-G450-business-jet-passengers.jpg"
  _snake_case							:	Any			=						Image.open(requests.get(snake_case__	,     stream=snake_case__					).raw					)
  return im


@torch.no_grad()
def       UpperCAmelCase__ (snake_case__							:      List[Any]	,     snake_case__							:      Optional[int]	,     snake_case__							:      List[str]	,     snake_case__							:      Tuple					):
  """simple docstring"""

  _snake_case							:	int			=						get_mobilevitva_config(snake_case__	,     snake_case__					)

  # load original state_dict
  _snake_case							:	Optional[int]			=						torch.load(snake_case__	,     map_location="""cpu"""					)

  # load huggingface model
  if task_name.startswith("""ade20k_"""					) or task_name.startswith("""voc_"""					):
    _snake_case							:	Any			=						MobileViTVaForSemanticSegmentation(snake_case__					).eval()
    _snake_case							:	List[Any]			=						False
  else:
    _snake_case							:	List[Any]			=						MobileViTVaForImageClassification(snake_case__					).eval()
    _snake_case							:	Optional[Any]			=						False

  # remove and rename some keys of load the original model
  _snake_case							:	Union[str, Any]			=						checkpoint
  remove_unused_keys(snake_case__					)
  _snake_case							:	List[str]			=						create_rename_keys(snake_case__	,     base_model=snake_case__					)
  for rename_key_src, rename_key_dest in rename_keys:
    rename_key(snake_case__	,     snake_case__	,     snake_case__					)

  # load modified state_dict
  model.load_state_dict(snake_case__					)

  # Check outputs on an image, prepared by MobileViTImageProcessor
  _snake_case							:	Optional[int]			=						MobileViTImageProcessor(crop_size=config.image_size	,     size=config.image_size + 32					)
  _snake_case							:	List[str]			=						image_processor(images=prepare_img()	,     return_tensors="""pt"""					)
  _snake_case							:	Optional[Any]			=						model(**snake_case__					)

  # verify classification model
  if task_name.startswith("""imagenet"""					):
    _snake_case							:	List[str]			=						outputs.logits
    _snake_case							:	Any			=						logits.argmax(-1					).item()
    print("""Predicted class:"""	,     model.config.idalabel[predicted_class_idx]					)
    if task_name.startswith("""imagenet1k_256"""					) and config.width_multiplier == 1.0:
      # expected_logits for base variant
      _snake_case							:	List[str]			=						torch.tensor([-1.6_3_3_6e0_0, -7.3_2_0_4e-0_2, -5.1_8_8_3e-0_1]					)
      assert torch.allclose(logits[0, :3]	,     snake_case__	,     atol=1e-4					)

  Path(snake_case__					).mkdir(exist_ok=snake_case__					)
  print(F"Saving model {task_name} to {pytorch_dump_folder_path}"					)
  model.save_pretrained(snake_case__					)
  print(F"Saving image processor to {pytorch_dump_folder_path}"					)
  image_processor.save_pretrained(snake_case__					)


if __name__ == "__main__":
     A_      			=       argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--task''',
         default='''imagenet1k_256''',
         type=str,
         help=(
             '''Name of the task for which the MobileViTV2 model you\'d like to convert is trained on . '''
             '''
                Classification (ImageNet-1k)
                    - MobileViTV2 (256x256) : imagenet1k_256
                    - MobileViTV2 (Trained on 256x256 and Finetuned on 384x384) : imagenet1k_384
                    - MobileViTV2 (Trained on ImageNet-21k and Finetuned on ImageNet-1k 256x256) :
                      imagenet21k_to_1k_256
                    - MobileViTV2 (Trained on ImageNet-21k, Finetuned on ImageNet-1k 256x256, and Finetuned on
                      ImageNet-1k 384x384) : imagenet21k_to_1k_384
                Segmentation
                    - ADE20K Dataset : ade20k_deeplabv3
                    - Pascal VOC 2012 Dataset: voc_deeplabv3
            '''
         ),
         choices=[
             '''imagenet1k_256''',
             '''imagenet1k_384''',
             '''imagenet21k_to_1k_256''',
             '''imagenet21k_to_1k_384''',
             '''ade20k_deeplabv3''',
             '''voc_deeplabv3''',
         ],
     )

     parser.add_argument(
         '''--orig_checkpoint_path''', required=True, type=str, help='''Path to the original state dict (.pt file).'''
     )
     parser.add_argument('''--orig_config_path''', required=True, type=str, help='''Path to the original config file.''')
     parser.add_argument(
         '''--pytorch_dump_folder_path''', required=True, type=str, help='''Path to the output PyTorch model directory.'''
     )

     A_      			=       parser.parse_args()
     convert_mobilevitva_checkpoint(
         args.task, args.orig_checkpoint_path, args.orig_config_path, args.pytorch_dump_folder_path
     )


