



'''simple docstring'''


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







	def   _a   (    self						)    ->    List[str]:
							'''simple docstring'''






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

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

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







	def   _a   (    self ,				**_lowerCAmelCase						)    ->    Tuple:
							'''simple docstring'''






							lowercase        =							self.add_kwargs_tokens_map.copy()
							kwargs.update(_lowerCAmelCase						)
							return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname ,				**_lowerCAmelCase						)







	def   _a   (    self ,				**_lowerCAmelCase						)    ->    Any:
							'''simple docstring'''






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







	def   _a   (    self ,				**_lowerCAmelCase						)    ->    List[str]:
							'''simple docstring'''






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







	def   _a   (    self						)    ->    int:
							'''simple docstring'''






							shutil.rmtree(self.tmpdirname						)







	def   _a   (    self						)    ->    List[str]:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase ,				feature_extractor=_lowerCAmelCase ,				decoder=_lowerCAmelCase						)

							processor.save_pretrained(self.tmpdirname						)
							lowercase        =							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   _a   (    self						)    ->    Optional[int]:
							'''simple docstring'''






							lowercase        =							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
							lowercase        =							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   _a   (    self						)    ->    Dict:
							'''simple docstring'''






							lowercase        =							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   _a   (    self						)    ->    Tuple:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase ,				feature_extractor=_lowerCAmelCase ,				decoder=_lowerCAmelCase						)

							lowercase        =							floats_list((3, 1000)						)

							lowercase        =							feature_extractor(_lowerCAmelCase ,				return_tensors="""np"""						)
							lowercase        =							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   _a   (    self						)    ->    int:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase ,				feature_extractor=_lowerCAmelCase ,				decoder=_lowerCAmelCase						)

							lowercase        =							"""This is a test string"""

							lowercase        =							processor(text=_lowerCAmelCase						)

							lowercase        =							tokenizer(_lowerCAmelCase						)

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







	def   _a   (    self ,				_lowerCAmelCase=(2, 10, 16) ,				_lowerCAmelCase=77						)    ->    List[Any]:
							'''simple docstring'''






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







	def   _a   (    self						)    ->    List[str]:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase ,				feature_extractor=_lowerCAmelCase ,				decoder=_lowerCAmelCase						)

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

							lowercase        =							processor.decode(_lowerCAmelCase						)

							lowercase        =							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   _a   (    self ,				_lowerCAmelCase						)    ->    str:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase ,				feature_extractor=_lowerCAmelCase ,				decoder=_lowerCAmelCase						)

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

							lowercase        =							list(_lowerCAmelCase						)

							with get_context("""fork"""						).Pool() as p:
													lowercase        =							decoder.decode_beams_batch(_lowerCAmelCase ,				_lowerCAmelCase						)

							lowercase     ,   lowercase     ,   lowercase        =							[], [], []
							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   _a   (    self						)    ->    int:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase ,				feature_extractor=_lowerCAmelCase ,				decoder=_lowerCAmelCase						)

							lowercase        =							self._get_dummy_logits()

							lowercase        =							15
							lowercase        =							-20.0
							lowercase        =							-4.0

							lowercase        =							processor.batch_decode(
							    _lowerCAmelCase ,				beam_width=_lowerCAmelCase ,				beam_prune_logp=_lowerCAmelCase ,				token_min_logp=_lowerCAmelCase ,				)
							lowercase        =							decoded_processor_out.text

							lowercase        =							list(_lowerCAmelCase						)

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

							lowercase        =							[d[0][0] for d in decoded_decoder_out]
							lowercase        =							[d[0][2] for d in decoded_decoder_out]
							lowercase        =							[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.9474] ,				_lowerCAmelCase ,				atol=1E-3						)						)







	def   _a   (    self						)    ->    Optional[Any]:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase ,				feature_extractor=_lowerCAmelCase ,				decoder=_lowerCAmelCase						)

							lowercase        =							self._get_dummy_logits()

							lowercase        =							2.0
							lowercase        =							5.0
							lowercase        =							-20.0
							lowercase        =							True

							lowercase        =							processor.batch_decode(
							    _lowerCAmelCase ,				alpha=_lowerCAmelCase ,				beta=_lowerCAmelCase ,				unk_score_offset=_lowerCAmelCase ,				lm_score_boundary=_lowerCAmelCase ,				)
							lowercase        =							decoded_processor_out.text

							lowercase        =							list(_lowerCAmelCase						)
							decoder.reset_params(
							    alpha=_lowerCAmelCase ,				beta=_lowerCAmelCase ,				unk_score_offset=_lowerCAmelCase ,				lm_score_boundary=_lowerCAmelCase ,				)

							with get_context("""fork"""						).Pool() as pool:
													lowercase        =							decoder.decode_beams_batch(
													    _lowerCAmelCase ,				_lowerCAmelCase ,				)

							lowercase        =							[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						)
							lowercase        =							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   _a   (    self						)    ->    List[str]:
							'''simple docstring'''






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

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

							lowercase        =							os.listdir(_lowerCAmelCase						)
							lowercase        =							["""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   _a   (    self						)    ->    Optional[int]:
							'''simple docstring'''






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

							lowercase        =							WavaVecaProcessorWithLM.from_pretrained(_lowerCAmelCase						)

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

							lowercase        =							os.listdir(_lowerCAmelCase						)
							lowercase        =							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   _a   (    self						)    ->    Any:
							'''simple docstring'''






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

							lowercase        =							floats_list((3, 1000)						)

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

							lowercase        =							self._get_dummy_logits()

							lowercase        =							processor_wavaveca.batch_decode(_lowerCAmelCase						)
							lowercase        =							processor_auto.batch_decode(_lowerCAmelCase						)

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







	def   _a   (    self						)    ->    List[Any]:
							'''simple docstring'''






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

							lowercase        =							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   _a   (    _lowerCAmelCase ,				_lowerCAmelCase						)    ->    Tuple:
							'''simple docstring'''






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







	def   _a   (    self						)    ->    str:
							'''simple docstring'''






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

							lowercase        =							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   _a   (    self						)    ->    Optional[Any]:
							'''simple docstring'''






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

							lowercase        =							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   _a   (    self						)    ->    Optional[int]:
							'''simple docstring'''






							import torch

							lowercase        =							load_dataset("""common_voice""" ,				"""en""" ,				split="""train""" ,				streaming=_lowerCAmelCase						)
							lowercase        =							ds.cast_column("""audio""" ,				datasets.Audio(sampling_rate=1_6000						)						)
							lowercase        =							iter(_lowerCAmelCase						)
							lowercase        =							next(_lowerCAmelCase						)

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

							with torch.no_grad():
													lowercase        =							model(_lowerCAmelCase						).logits.cpu().numpy()

							lowercase        =							processor.decode(logits[0] ,				output_word_offsets=_lowerCAmelCase						)

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

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

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

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







'''simple docstring'''


import argparse
import json
from pathlib import Path

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

from transformers import (
    MobileViTConfig,
    MobileViTForImageClassification,
    MobileViTForSemanticSegmentation,
    MobileViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowercase_      :				List[str]   					=    logging.get_logger(__name__)







def      SCREAMING_SNAKE_CASE				(				lowercase_    :      int						):
						lowercase        =							MobileViTConfig()

						# size of the architecture
						if "mobilevit_s" in mobilevit_name:
												lowercase        =							[144, 192, 240]
												lowercase        =							[16, 32, 64, 96, 128, 160, 640]
						elif "mobilevit_xs" in mobilevit_name:
												lowercase        =							[96, 120, 144]
												lowercase        =							[16, 32, 48, 64, 80, 96, 384]
						elif "mobilevit_xxs" in mobilevit_name:
												lowercase        =							[64, 80, 96]
												lowercase        =							[16, 16, 24, 48, 64, 80, 320]
												lowercase        =							0.05
												lowercase        =							2.0

						if mobilevit_name.startswith("""deeplabv3_"""						):
												lowercase        =							512
												lowercase        =							16
												lowercase        =							21
												lowercase        =							"""pascal-voc-id2label.json"""
						else:
												lowercase        =							1000
												lowercase        =							"""imagenet-1k-id2label.json"""

						lowercase        =							"""huggingface/label-files"""
						lowercase        =							json.load(open(hf_hub_download(lowercase_ ,  lowercase_ ,  repo_type="""dataset"""						) ,  """r"""						)						)
						lowercase        =							{int(lowercase_						): v for k, v in idalabel.items()}
						lowercase        =							idalabel
						lowercase        =							{v: k for k, v in idalabel.items()}

						return config







def      SCREAMING_SNAKE_CASE				(				lowercase_    :      Any ,  lowercase_    :      Any=False						):
						for i in range(1 ,  6						):
												if F"""layer_{i}.""" in name:
																		lowercase        =							name.replace(F"""layer_{i}.""" ,  F"""encoder.layer.{i - 1}."""						)

						if "conv_1." in name:
												lowercase        =							name.replace("""conv_1.""" ,  """conv_stem."""						)
						if ".block." in name:
												lowercase        =							name.replace(""".block.""" ,  """."""						)

						if "exp_1x1" in name:
												lowercase        =							name.replace("""exp_1x1""" ,  """expand_1x1"""						)
						if "red_1x1" in name:
												lowercase        =							name.replace("""red_1x1""" ,  """reduce_1x1"""						)
						if ".local_rep.conv_3x3." in name:
												lowercase        =							name.replace(""".local_rep.conv_3x3.""" ,  """.conv_kxk."""						)
						if ".local_rep.conv_1x1." in name:
												lowercase        =							name.replace(""".local_rep.conv_1x1.""" ,  """.conv_1x1."""						)
						if ".norm." in name:
												lowercase        =							name.replace(""".norm.""" ,  """.normalization."""						)
						if ".conv." in name:
												lowercase        =							name.replace(""".conv.""" ,  """.convolution."""						)
						if ".conv_proj." in name:
												lowercase        =							name.replace(""".conv_proj.""" ,  """.conv_projection."""						)

						for i in range(0 ,  2						):
												for j in range(0 ,  4						):
																		if F""".{i}.{j}.""" in name:
																								lowercase        =							name.replace(F""".{i}.{j}.""" ,  F""".{i}.layer.{j}."""						)

						for i in range(2 ,  6						):
												for j in range(0 ,  4						):
																		if F""".{i}.{j}.""" in name:
																								lowercase        =							name.replace(F""".{i}.{j}.""" ,  F""".{i}."""						)
																								if "expand_1x1" in name:
																														lowercase        =							name.replace("""expand_1x1""" ,  """downsampling_layer.expand_1x1"""						)
																								if "conv_3x3" in name:
																														lowercase        =							name.replace("""conv_3x3""" ,  """downsampling_layer.conv_3x3"""						)
																								if "reduce_1x1" in name:
																														lowercase        =							name.replace("""reduce_1x1""" ,  """downsampling_layer.reduce_1x1"""						)

						for i in range(2 ,  5						):
												if F""".global_rep.{i}.weight""" in name:
																		lowercase        =							name.replace(F""".global_rep.{i}.weight""" ,  """.layernorm.weight"""						)
												if F""".global_rep.{i}.bias""" in name:
																		lowercase        =							name.replace(F""".global_rep.{i}.bias""" ,  """.layernorm.bias"""						)

						if ".global_rep." in name:
												lowercase        =							name.replace(""".global_rep.""" ,  """.transformer."""						)
						if ".pre_norm_mha.0." in name:
												lowercase        =							name.replace(""".pre_norm_mha.0.""" ,  """.layernorm_before."""						)
						if ".pre_norm_mha.1.out_proj." in name:
												lowercase        =							name.replace(""".pre_norm_mha.1.out_proj.""" ,  """.attention.output.dense."""						)
						if ".pre_norm_ffn.0." in name:
												lowercase        =							name.replace(""".pre_norm_ffn.0.""" ,  """.layernorm_after."""						)
						if ".pre_norm_ffn.1." in name:
												lowercase        =							name.replace(""".pre_norm_ffn.1.""" ,  """.intermediate.dense."""						)
						if ".pre_norm_ffn.4." in name:
												lowercase        =							name.replace(""".pre_norm_ffn.4.""" ,  """.output.dense."""						)
						if ".transformer." in name:
												lowercase        =							name.replace(""".transformer.""" ,  """.transformer.layer."""						)

						if ".aspp_layer." in name:
												lowercase        =							name.replace(""".aspp_layer.""" ,  """."""						)
						if ".aspp_pool." in name:
												lowercase        =							name.replace(""".aspp_pool.""" ,  """."""						)
						if "seg_head." in name:
												lowercase        =							name.replace("""seg_head.""" ,  """segmentation_head."""						)
						if "segmentation_head.classifier.classifier." in name:
												lowercase        =							name.replace("""segmentation_head.classifier.classifier.""" ,  """segmentation_head.classifier."""						)

						if "classifier.fc." in name:
												lowercase        =							name.replace("""classifier.fc.""" ,  """classifier."""						)
						elif (not base_model) and ("segmentation_head." not in name):
												lowercase        =							"""mobilevit.""" + name

						return name







def      SCREAMING_SNAKE_CASE				(				lowercase_    :      Optional[Any] ,  lowercase_    :      List[Any] ,  lowercase_    :      str=False						):
						if base_model:
												lowercase        =							""""""
						else:
												lowercase        =							"""mobilevit."""

						for key in orig_state_dict.copy().keys():
												lowercase        =							orig_state_dict.pop(lowercase_						)

												if key[:8] == "encoder.":
																		lowercase        =							key[8:]

												if "qkv" in key:
																		lowercase        =							key.split("""."""						)
																		lowercase        =							int(key_split[0][6:]						) - 1
																		lowercase        =							int(key_split[3]						)
																		lowercase        =							model.get_submodule(F"""{model_prefix}encoder.layer.{layer_num}"""						)
																		lowercase        =							layer.transformer.layer[transformer_num].attention.attention.all_head_size
																		lowercase        =							(
																		    F"""{model_prefix}encoder.layer.{layer_num}.transformer.layer.{transformer_num}.attention.attention."""
																		)
																		if "weight" in key:
																								lowercase        =							val[:dim, :]
																								lowercase        =							val[dim : dim * 2, :]
																								lowercase        =							val[-dim:, :]
																		else:
																								lowercase        =							val[:dim]
																								lowercase        =							val[dim : dim * 2]
																								lowercase        =							val[-dim:]
												else:
																		lowercase        =							val

						return orig_state_dict







def      SCREAMING_SNAKE_CASE				(				):
						lowercase        =							"""http://images.cocodataset.org/val2017/000000039769.jpg"""
						lowercase        =							Image.open(requests.get(lowercase_ ,  stream=lowercase_						).raw						)
						return im







@torch.no_grad()
def      SCREAMING_SNAKE_CASE				(				lowercase_    :      Dict ,  lowercase_    :      List[Any] ,  lowercase_    :      Any ,  lowercase_    :      List[str]=False						):
						lowercase        =							get_mobilevit_config(lowercase_						)

						# load original state_dict
						lowercase        =							torch.load(lowercase_ ,  map_location="""cpu"""						)

						# load 🤗 model
						if mobilevit_name.startswith("""deeplabv3_"""						):
												lowercase        =							MobileViTForSemanticSegmentation(lowercase_						).eval()
						else:
												lowercase        =							MobileViTForImageClassification(lowercase_						).eval()

						lowercase        =							convert_state_dict(lowercase_ ,  lowercase_						)
						model.load_state_dict(lowercase_						)

						# Check outputs on an image, prepared by MobileViTImageProcessor
						lowercase        =							MobileViTImageProcessor(crop_size=config.image_size ,  size=config.image_size + 32						)
						lowercase        =							image_processor(images=prepare_img() ,  return_tensors="""pt"""						)
						lowercase        =							model(**lowercase_						)
						lowercase        =							outputs.logits

						if mobilevit_name.startswith("""deeplabv3_"""						):
												assert logits.shape == (1, 21, 32, 32)

												if mobilevit_name == "deeplabv3_mobilevit_s":
																		lowercase        =							torch.tensor(
																		    [
																		        [[6.2_065, 6.1_292, 6.2_070], [6.1_079, 6.1_254, 6.1_747], [6.0_042, 6.1_071, 6.1_034]],
																		        [[-6.9_253, -6.8_653, -7.0_398], [-7.3_218, -7.3_983, -7.3_670], [-7.1_961, -7.2_482, -7.1_569]],
																		        [[-4.4_723, -4.4_348, -4.3_769], [-5.3_629, -5.4_632, -5.4_598], [-5.1_587, -5.3_402, -5.5_059]],
																		    ]						)
												elif mobilevit_name == "deeplabv3_mobilevit_xs":
																		lowercase        =							torch.tensor(
																		    [
																		        [[5.4_449, 5.5_733, 5.6_314], [5.1_815, 5.3_930, 5.5_963], [5.1_656, 5.4_333, 5.4_853]],
																		        [[-9.4_423, -9.7_766, -9.6_714], [-9.1_581, -9.5_720, -9.5_519], [-9.1_006, -9.6_458, -9.5_703]],
																		        [[-7.7_721, -7.3_716, -7.1_583], [-8.4_599, -8.0_624, -7.7_944], [-8.4_172, -7.8_366, -7.5_025]],
																		    ]						)
												elif mobilevit_name == "deeplabv3_mobilevit_xxs":
																		lowercase        =							torch.tensor(
																		    [
																		        [[6.9_811, 6.9_743, 7.3_123], [7.1_777, 7.1_931, 7.3_938], [7.5_633, 7.8_050, 7.8_901]],
																		        [[-10.5_536, -10.2_332, -10.2_924], [-10.2_336, -9.8_624, -9.5_964], [-10.8_840, -10.8_158, -10.6_659]],
																		        [[-3.4_938, -3.0_631, -2.8_620], [-3.4_205, -2.8_135, -2.6_875], [-3.4_179, -2.7_945, -2.8_750]],
																		    ]						)
												else:
																		raise ValueError(F"""Unknown mobilevit_name: {mobilevit_name}"""						)

												assert torch.allclose(logits[0, :3, :3, :3] ,  lowercase_ ,  atol=1E-4						)
						else:
												assert logits.shape == (1, 1000)

												if mobilevit_name == "mobilevit_s":
																		lowercase        =							torch.tensor([-0.9_866, 0.2_392, -1.1_241]						)
												elif mobilevit_name == "mobilevit_xs":
																		lowercase        =							torch.tensor([-2.4_761, -0.9_399, -1.9_587]						)
												elif mobilevit_name == "mobilevit_xxs":
																		lowercase        =							torch.tensor([-1.9_364, -1.2_327, -0.4_653]						)
												else:
																		raise ValueError(F"""Unknown mobilevit_name: {mobilevit_name}"""						)

												assert torch.allclose(logits[0, :3] ,  lowercase_ ,  atol=1E-4						)

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

						if push_to_hub:
												lowercase        =							{
												    """mobilevit_s""": """mobilevit-small""",
												    """mobilevit_xs""": """mobilevit-x-small""",
												    """mobilevit_xxs""": """mobilevit-xx-small""",
												    """deeplabv3_mobilevit_s""": """deeplabv3-mobilevit-small""",
												    """deeplabv3_mobilevit_xs""": """deeplabv3-mobilevit-x-small""",
												    """deeplabv3_mobilevit_xxs""": """deeplabv3-mobilevit-xx-small""",
												}

												print("""Pushing to the hub..."""						)
												lowercase        =							model_mapping[mobilevit_name]
												image_processor.push_to_hub(lowercase_ ,  organization="""apple"""						)
												model.push_to_hub(lowercase_ ,  organization="""apple"""						)


if __name__ == "__main__":
						lowercase_      :				List[str]   					=    argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--mobilevit_name''',
						    default='''mobilevit_s''',
						    type=str,
						    help=(
						        '''Name of the MobileViT model you\'d like to convert. Should be one of \'mobilevit_s\', \'mobilevit_xs\','''
						        ''' \'mobilevit_xxs\', \'deeplabv3_mobilevit_s\', \'deeplabv3_mobilevit_xs\', \'deeplabv3_mobilevit_xxs\'.'''
						    ),
						)
						parser.add_argument(
						    '''--checkpoint_path''', required=True, type=str, help='''Path to the original state dict (.pt file).'''
						)
						parser.add_argument(
						    '''--pytorch_dump_folder_path''', required=True, type=str, help='''Path to the output PyTorch model directory.'''
						)
						parser.add_argument(
						    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
						)

						lowercase_      :				List[str]   					=    parser.parse_args()
						convert_movilevit_checkpoint(
						    args.mobilevit_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
						)


