






"""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      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			Dict:
				UpperCAmelCase_     : List[str]			     =					"| <pad> <unk> <s> </s> a b c d e f g h i j k".split()
				UpperCAmelCase_     : Optional[Any]			     =					dict(zip(lowerCAmelCase_	,	range(len(lowerCAmelCase_						)						)						)						)

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

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





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[Any]	,	**lowerCAmelCase_					: List[Any]						)						->			List[str]:
				UpperCAmelCase_     : List[str]			     =					self.add_kwargs_tokens_map.copy()
				kwargs.update(lowerCAmelCase_						)
				return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname	,	**lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Dict	,	**lowerCAmelCase_					: Tuple						)						->			Optional[int]:
				return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname	,	**lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple	,	**lowerCAmelCase_					: Optional[int]						)						->			int:
				return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name	,	**lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]						)						->			List[Any]:
				shutil.rmtree(self.tmpdirname						)





	def      _SCREAMING_SNAKE_CASE      (      self					: str						)						->			str:
				UpperCAmelCase_     : Union[str, Any]			     =					self.get_tokenizer()
				UpperCAmelCase_     : List[str]			     =					self.get_feature_extractor()
				UpperCAmelCase_     : List[str]			     =					self.get_decoder()

				UpperCAmelCase_     : Dict			     =					WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase_	,	feature_extractor=lowerCAmelCase_	,	decoder=lowerCAmelCase_						)

				processor.save_pretrained(self.tmpdirname						)
				UpperCAmelCase_     : Tuple			     =					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      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			str:
				UpperCAmelCase_     : Optional[Any]			     =					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_     : List[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      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			Any:
				UpperCAmelCase_     : Optional[int]			     =					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      _SCREAMING_SNAKE_CASE      (      self					: Optional[Any]						)						->			Any:
				UpperCAmelCase_     : List[str]			     =					self.get_feature_extractor()
				UpperCAmelCase_     : Dict			     =					self.get_tokenizer()
				UpperCAmelCase_     : Dict			     =					self.get_decoder()

				UpperCAmelCase_     : str			     =					WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase_	,	feature_extractor=lowerCAmelCase_	,	decoder=lowerCAmelCase_						)

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

				UpperCAmelCase_     : Optional[Any]			     =					feature_extractor(lowerCAmelCase_	,	return_tensors="np"						)
				UpperCAmelCase_     : Tuple			     =					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      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			Optional[int]:
				UpperCAmelCase_     : str			     =					self.get_feature_extractor()
				UpperCAmelCase_     : Optional[int]			     =					self.get_tokenizer()
				UpperCAmelCase_     : List[str]			     =					self.get_decoder()

				UpperCAmelCase_     : List[str]			     =					WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase_	,	feature_extractor=lowerCAmelCase_	,	decoder=lowerCAmelCase_						)

				UpperCAmelCase_     : int			     =					"This is a test string"

				UpperCAmelCase_     : Union[str, Any]			     =					processor(text=lowerCAmelCase_						)

				UpperCAmelCase_     : Union[str, Any]			     =					tokenizer(lowerCAmelCase_						)

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





	def      _SCREAMING_SNAKE_CASE      (      self					: str	,	lowerCAmelCase_					: int=(2, 10, 16)	,	lowerCAmelCase_					: List[str]=77						)						->			Dict:
				np.random.seed(lowerCAmelCase_						)
				return np.random.rand(*lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Any						)						->			Tuple:
				UpperCAmelCase_     : List[str]			     =					self.get_feature_extractor()
				UpperCAmelCase_     : List[str]			     =					self.get_tokenizer()
				UpperCAmelCase_     : Dict			     =					self.get_decoder()

				UpperCAmelCase_     : List[str]			     =					WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase_	,	feature_extractor=lowerCAmelCase_	,	decoder=lowerCAmelCase_						)

				UpperCAmelCase_     : Union[str, Any]			     =					self._get_dummy_logits(shape=(10, 16)	,	seed=13						)

				UpperCAmelCase_     : Tuple			     =					processor.decode(lowerCAmelCase_						)

				UpperCAmelCase_     : int			     =					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      _SCREAMING_SNAKE_CASE      (      self					: Dict	,	lowerCAmelCase_					: Tuple						)						->			Optional[int]:
				UpperCAmelCase_     : Optional[int]			     =					self.get_feature_extractor()
				UpperCAmelCase_     : List[str]			     =					self.get_tokenizer()
				UpperCAmelCase_     : int			     =					self.get_decoder()

				UpperCAmelCase_     : Optional[Any]			     =					WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase_	,	feature_extractor=lowerCAmelCase_	,	decoder=lowerCAmelCase_						)

				UpperCAmelCase_     : str			     =					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_     : str			     =					processor.batch_decode(lowerCAmelCase_						)
				else:
							with get_context(lowerCAmelCase_						).Pool() as pool:
										UpperCAmelCase_     : Optional[int]			     =					processor.batch_decode(lowerCAmelCase_	,	lowerCAmelCase_						)

				UpperCAmelCase_     : Dict			     =					list(lowerCAmelCase_						)

				with get_context("fork"						).Pool() as p:
							UpperCAmelCase_     : Union[str, Any]			     =					decoder.decode_beams_batch(lowerCAmelCase_	,	lowerCAmelCase_						)

				UpperCAmelCase_     ,	UpperCAmelCase_     ,	UpperCAmelCase_     : Optional[Any]			     =					[], [], []
				for beams in decoded_beams:
							texts_decoder.append(beams[0][0]						)
							logit_scores_decoder.append(beams[0][-2]						)
							lm_scores_decoder.append(beams[0][-1]						)

				self.assertListEqual(lowerCAmelCase_	,	decoded_processor.text						)
				self.assertListEqual(["<s> <s> </s>", "<s> <s> <s>"]	,	decoded_processor.text						)
				self.assertListEqual(lowerCAmelCase_	,	decoded_processor.logit_score						)
				self.assertListEqual(lowerCAmelCase_	,	decoded_processor.lm_score						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Union[str, Any]						)						->			int:
				UpperCAmelCase_     : List[Any]			     =					self.get_feature_extractor()
				UpperCAmelCase_     : Union[str, Any]			     =					self.get_tokenizer()
				UpperCAmelCase_     : Union[str, Any]			     =					self.get_decoder()

				UpperCAmelCase_     : List[Any]			     =					WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase_	,	feature_extractor=lowerCAmelCase_	,	decoder=lowerCAmelCase_						)

				UpperCAmelCase_     : str			     =					self._get_dummy_logits()

				UpperCAmelCase_     : Tuple			     =					15
				UpperCAmelCase_     : List[str]			     =					-2_0.0
				UpperCAmelCase_     : int			     =					-4.0

				UpperCAmelCase_     : Union[str, Any]			     =					processor.batch_decode(
				    lowerCAmelCase_	,	beam_width=lowerCAmelCase_	,	beam_prune_logp=lowerCAmelCase_	,	token_min_logp=lowerCAmelCase_	,	)
				UpperCAmelCase_     : Union[str, Any]			     =					decoded_processor_out.text

				UpperCAmelCase_     : int			     =					list(lowerCAmelCase_						)

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

				UpperCAmelCase_     : int			     =					[d[0][0] for d in decoded_decoder_out]
				UpperCAmelCase_     : int			     =					[d[0][2] for d in decoded_decoder_out]
				UpperCAmelCase_     : Tuple			     =					[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([-2_0.0_5_4, -1_8.4_4_7]	,	lowerCAmelCase_	,	atol=1e-3						)						)

				self.assertTrue(np.array_equal(lowerCAmelCase_	,	decoded_processor_out.lm_score						)						)
				self.assertTrue(np.allclose([-1_5.5_5_4, -1_3.9_4_7_4]	,	lowerCAmelCase_	,	atol=1e-3						)						)





	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]						)						->			str:
				UpperCAmelCase_     : int			     =					self.get_feature_extractor()
				UpperCAmelCase_     : Union[str, Any]			     =					self.get_tokenizer()
				UpperCAmelCase_     : Tuple			     =					self.get_decoder()

				UpperCAmelCase_     : int			     =					WavaVecaProcessorWithLM(tokenizer=lowerCAmelCase_	,	feature_extractor=lowerCAmelCase_	,	decoder=lowerCAmelCase_						)

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

				UpperCAmelCase_     : int			     =					2.0
				UpperCAmelCase_     : Tuple			     =					5.0
				UpperCAmelCase_     : Tuple			     =					-2_0.0
				UpperCAmelCase_     : Union[str, Any]			     =					True

				UpperCAmelCase_     : Optional[Any]			     =					processor.batch_decode(
				    lowerCAmelCase_	,	alpha=lowerCAmelCase_	,	beta=lowerCAmelCase_	,	unk_score_offset=lowerCAmelCase_	,	lm_score_boundary=lowerCAmelCase_	,	)
				UpperCAmelCase_     : Tuple			     =					decoded_processor_out.text

				UpperCAmelCase_     : Union[str, Any]			     =					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_     : int			     =					decoder.decode_beams_batch(
							    lowerCAmelCase_	,	lowerCAmelCase_	,	)

				UpperCAmelCase_     : Any			     =					[d[0][0] for d in decoded_decoder_out]

				self.assertListEqual(lowerCAmelCase_	,	lowerCAmelCase_						)
				self.assertListEqual(["<s> </s> <s> </s> </s>", "</s> </s> <s> </s> </s>"]	,	lowerCAmelCase_						)
				UpperCAmelCase_     : 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	,	-2_0.0						)
				self.assertEqual(lm_model.score_boundary	,	lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			str:
				UpperCAmelCase_     : List[Any]			     =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"						)

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

				UpperCAmelCase_     : str			     =					os.listdir(lowerCAmelCase_						)
				UpperCAmelCase_     : Dict			     =					["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      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			Tuple:
				UpperCAmelCase_     : str			     =					snapshot_download("hf-internal-testing/processor_with_lm"						)

				UpperCAmelCase_     : Optional[Any]			     =					WavaVecaProcessorWithLM.from_pretrained(lowerCAmelCase_						)

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

				UpperCAmelCase_     : str			     =					os.listdir(lowerCAmelCase_						)
				UpperCAmelCase_     : List[str]			     =					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      _SCREAMING_SNAKE_CASE      (      self					: Optional[Any]						)						->			int:
				UpperCAmelCase_     : Union[str, Any]			     =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"						)
				UpperCAmelCase_     : List[str]			     =					AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"						)

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

				UpperCAmelCase_     : Optional[Any]			     =					processor_wavaveca(lowerCAmelCase_	,	return_tensors="np"						)
				UpperCAmelCase_     : Union[str, Any]			     =					processor_auto(lowerCAmelCase_	,	return_tensors="np"						)

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

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

				UpperCAmelCase_     : List[str]			     =					processor_wavaveca.batch_decode(lowerCAmelCase_						)
				UpperCAmelCase_     : int			     =					processor_auto.batch_decode(lowerCAmelCase_						)

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





	def      _SCREAMING_SNAKE_CASE      (      self					: Union[str, Any]						)						->			Any:
				UpperCAmelCase_     : Dict			     =					self.get_feature_extractor()
				UpperCAmelCase_     : Any			     =					self.get_tokenizer()
				UpperCAmelCase_     : List[Any]			     =					self.get_decoder()

				UpperCAmelCase_     : Optional[int]			     =					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      _SCREAMING_SNAKE_CASE      (      lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Tuple						)						->			List[str]:
				UpperCAmelCase_     : List[Any]			     =					[d[key] for d in offsets]
				return retrieved_list





	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			str:
				UpperCAmelCase_     : Tuple			     =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"						)
				UpperCAmelCase_     : Optional[Any]			     =					self._get_dummy_logits()[0]

				UpperCAmelCase_     : List[str]			     =					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      _SCREAMING_SNAKE_CASE      (      self					: List[Any]						)						->			Any:
				UpperCAmelCase_     : int			     =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"						)
				UpperCAmelCase_     : Any			     =					self._get_dummy_logits()

				UpperCAmelCase_     : Union[str, Any]			     =					processor.batch_decode(lowerCAmelCase_	,	output_word_offsets=lowerCAmelCase_						)

				# check Wav2Vec2CTCTokenizerOutput keys for word
				self.assertEqual(len(outputs.keys()						)	,	4						)
				self.assertTrue("text" in outputs						)
				self.assertTrue("word_offsets" in outputs						)
				self.assertTrue(isinstance(lowerCAmelCase_	,	lowerCAmelCase_						)						)

				self.assertListEqual(
				    [" ".join(self.get_from_offsets(lowerCAmelCase_	,	"word"						)						) for o in outputs["word_offsets"]]	,	outputs.text						)
				self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]	,	"word"						)	,	["<s>", "<s>", "</s>"]						)
				self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]	,	"start_offset"						)	,	[0, 2, 4]						)
				self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]	,	"end_offset"						)	,	[1, 3, 5]						)





	@slow
	@require_torch
	@require_torchaudio
	def      _SCREAMING_SNAKE_CASE      (      self					: Any						)						->			Union[str, Any]:
				import torch

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

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

				with torch.no_grad():
							UpperCAmelCase_     : Optional[Any]			     =					model(lowerCAmelCase_						).logits.cpu().numpy()

				UpperCAmelCase_     : Any			     =					processor.decode(logits[0]	,	output_word_offsets=lowerCAmelCase_						)

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

				UpperCAmelCase_     : Optional[int]			     =					"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_     : Tuple			     =					torch.tensor(self.get_from_offsets(lowerCAmelCase_	,	"start_time"						)						)
				UpperCAmelCase_     : List[Any]			     =					torch.tensor(self.get_from_offsets(lowerCAmelCase_	,	"end_time"						)						)

				# fmt: off
				UpperCAmelCase_     : List[Any]			     =					torch.tensor([1.4_1_9_9, 1.6_5_9_9, 2.2_5_9_9, 3.0, 3.2_4, 3.5_9_9_9, 3.7_9_9_9, 4.0_9_9_9, 4.2_6, 4.9_4, 5.2_8, 5.6_5_9_9, 5.7_8, 5.9_4, 6.3_2, 6.5_3_9_9, 6.6_5_9_9]						)
				UpperCAmelCase_     : Union[str, Any]			     =					torch.tensor([1.5_3_9_9, 1.8_9_9_9, 2.9, 3.1_6, 3.5_3_9_9, 3.7_2, 4.0_1_9_9, 4.1_7_9_9, 4.7_6, 5.1_5_9_9, 5.5_5_9_9, 5.6_9_9_9, 5.8_6, 6.1_9_9_9, 6.3_8, 6.6_1_9_9, 6.9_4]						)
				# fmt: on

				self.assertTrue(torch.allclose(lowerCAmelCase_	,	lowerCAmelCase_	,	atol=0.0_1						)						)
				self.assertTrue(torch.allclose(lowerCAmelCase_	,	lowerCAmelCase_	,	atol=0.0_1						)						)











"""simple docstring"""
import argparse

import requests
import torch
from PIL import Image

from transformers import SwinConfig, SwinForMaskedImageModeling, ViTImageProcessor

def 				snake_case      (   A__						):
			UpperCAmelCase_     : Dict			     =					SwinConfig(image_size=1_92						)

			if "base" in model_name:
						UpperCAmelCase_     : Any			     =					6
						UpperCAmelCase_     : Optional[Any]			     =					1_28
						UpperCAmelCase_     : Optional[int]			     =					(2, 2, 18, 2)
						UpperCAmelCase_     : List[str]			     =					(4, 8, 16, 32)
			elif "large" in model_name:
						UpperCAmelCase_     : Dict			     =					12
						UpperCAmelCase_     : int			     =					1_92
						UpperCAmelCase_     : List[Any]			     =					(2, 2, 18, 2)
						UpperCAmelCase_     : int			     =					(6, 12, 24, 48)
			else:
						raise ValueError("Model not supported, only supports base and large variants"						)

			UpperCAmelCase_     : str			     =					window_size
			UpperCAmelCase_     : Any			     =					embed_dim
			UpperCAmelCase_     : int			     =					depths
			UpperCAmelCase_     : Any			     =					num_heads

			return config

def 				snake_case      (   A__						):
			if "encoder.mask_token" in name:
						UpperCAmelCase_     : str			     =					name.replace("encoder.mask_token"	,"embeddings.mask_token"						)
			if "encoder.patch_embed.proj" in name:
						UpperCAmelCase_     : Optional[int]			     =					name.replace("encoder.patch_embed.proj"	,"embeddings.patch_embeddings.projection"						)
			if "encoder.patch_embed.norm" in name:
						UpperCAmelCase_     : List[str]			     =					name.replace("encoder.patch_embed.norm"	,"embeddings.norm"						)
			if "attn.proj" in name:
						UpperCAmelCase_     : Optional[Any]			     =					name.replace("attn.proj"	,"attention.output.dense"						)
			if "attn" in name:
						UpperCAmelCase_     : Any			     =					name.replace("attn"	,"attention.self"						)
			if "norm1" in name:
						UpperCAmelCase_     : str			     =					name.replace("norm1"	,"layernorm_before"						)
			if "norm2" in name:
						UpperCAmelCase_     : Tuple			     =					name.replace("norm2"	,"layernorm_after"						)
			if "mlp.fc1" in name:
						UpperCAmelCase_     : List[str]			     =					name.replace("mlp.fc1"	,"intermediate.dense"						)
			if "mlp.fc2" in name:
						UpperCAmelCase_     : str			     =					name.replace("mlp.fc2"	,"output.dense"						)

			if name == "encoder.norm.weight":
						UpperCAmelCase_     : List[str]			     =					"layernorm.weight"
			if name == "encoder.norm.bias":
						UpperCAmelCase_     : int			     =					"layernorm.bias"

			if "decoder" in name:
						pass
			else:
						UpperCAmelCase_     : Any			     =					"swin." + name

			return name

def 				snake_case      (   A__	,A__						):
			for key in orig_state_dict.copy().keys():
						UpperCAmelCase_     : Tuple			     =					orig_state_dict.pop(A__						)

						if "attn_mask" in key:
									pass
						elif "qkv" in key:
									UpperCAmelCase_     : Optional[int]			     =					key.split("."						)
									UpperCAmelCase_     : str			     =					int(key_split[2]						)
									UpperCAmelCase_     : Union[str, Any]			     =					int(key_split[4]						)
									UpperCAmelCase_     : Optional[Any]			     =					model.swin.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

									if "weight" in key:
												UpperCAmelCase_     : List[Any]			     =					val[:dim, :]
												UpperCAmelCase_     : str			     =					val[
												    dim : dim * 2, :
												]
												UpperCAmelCase_     : str			     =					val[-dim:, :]
									else:
												UpperCAmelCase_     : List[str]			     =					val[
												    :dim
												]
												UpperCAmelCase_     : str			     =					val[
												    dim : dim * 2
												]
												UpperCAmelCase_     : Optional[Any]			     =					val[
												    -dim:
												]
						else:
									UpperCAmelCase_     : Tuple			     =					val

			return orig_state_dict

def 				snake_case      (   A__	,A__	,A__	,A__						):
			UpperCAmelCase_     : List[Any]			     =					torch.load(A__	,map_location="cpu"						)["model"]

			UpperCAmelCase_     : Optional[Any]			     =					get_swin_config(A__						)
			UpperCAmelCase_     : List[Any]			     =					SwinForMaskedImageModeling(A__						)
			model.eval()

			UpperCAmelCase_     : str			     =					convert_state_dict(A__	,A__						)
			model.load_state_dict(A__						)

			UpperCAmelCase_     : int			     =					"http://images.cocodataset.org/val2017/000000039769.jpg"

			UpperCAmelCase_     : int			     =					ViTImageProcessor(size={"height": 1_92, "width": 1_92}						)
			UpperCAmelCase_     : Any			     =					Image.open(requests.get(A__	,stream=A__						).raw						)
			UpperCAmelCase_     : Any			     =					image_processor(images=A__	,return_tensors="pt"						)

			with torch.no_grad():
						UpperCAmelCase_     : List[Any]			     =					model(**A__						).logits

			print(outputs.keys()						)
			print("Looks ok!"						)

			if pytorch_dump_folder_path is not None:
						print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""						)
						model.save_pretrained(A__						)

						print(F"""Saving image processor to {pytorch_dump_folder_path}"""						)
						image_processor.save_pretrained(A__						)

			if push_to_hub:
						print(F"""Pushing model and image processor for {model_name} to hub"""						)
						model.push_to_hub(F"""microsoft/{model_name}"""						)
						image_processor.push_to_hub(F"""microsoft/{model_name}"""						)


if __name__ == "__main__":
			lowerCamelCase_					=      argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '''--model_name''',
			    default='''swin-base-simmim-window6-192''',
			    type=str,
			    choices=['''swin-base-simmim-window6-192''', '''swin-large-simmim-window12-192'''],
			    help='''Name of the Swin SimMIM model you\'d like to convert.''',
			)
			parser.add_argument(
			    '''--checkpoint_path''',
			    default='''/Users/nielsrogge/Documents/SwinSimMIM/simmim_pretrain__swin_base__img192_window6__100ep.pth''',
			    type=str,
			    help='''Path to the original PyTorch checkpoint (.pth file).''',
			)
			parser.add_argument(
			    '''--pytorch_dump_folder_path''', default=None, 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.'''
			)

			lowerCamelCase_					=      parser.parse_args()
			convert_swin_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)



