





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






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									lowercase_	:   str  				=     '''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
									lowercase_	:   int  				=     dict(zip(UpperCamelCase__,   range(len(UpperCamelCase__  )  )  )  )

									lowercase_	:   Any  				=     {
									    '''unk_token''': '''<unk>''',
									    '''bos_token''': '''<s>''',
									    '''eos_token''': '''</s>''',
									}
									lowercase_	:   str  				=     {
									    '''feature_size''': 1,
									    '''padding_value''': 0.0,
									    '''sampling_rate''': 1_60_00,
									    '''return_attention_mask''': False,
									    '''do_normalize''': True,
									}

									lowercase_	:   Optional[Any]  				=     tempfile.mkdtemp()
									lowercase_	:   Dict  				=     os.path.join(self.tmpdirname,   VOCAB_FILES_NAMES["""vocab_file"""]  )
									lowercase_	:   Any  				=     os.path.join(self.tmpdirname,   UpperCamelCase__  )
									with open(self.vocab_file,   """w""",   encoding="""utf-8"""  ) as fp:
																fp.write(json.dumps(UpperCamelCase__  ) + """\n"""  )

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

									# load decoder from hub
									lowercase_	:   List[Any]  				=     '''hf-internal-testing/ngram-beam-search-decoder'''






		def 						snake_case__				(							self,   **snake_case__  )       -> Tuple:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     self.add_kwargs_tokens_map.copy()
									kwargs.update(UpperCamelCase__  )
									return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname,   **UpperCamelCase__  )






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






									"""simple docstring"""
									return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname,   **UpperCamelCase__  )






		def 						snake_case__				(							self,   **snake_case__  )       -> List[Any]:






									"""simple docstring"""
									return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name,   **UpperCamelCase__  )






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									shutil.rmtree(self.tmpdirname  )






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     self.get_tokenizer()
									lowercase_	:   Any  				=     self.get_feature_extractor()
									lowercase_	:   List[Any]  				=     self.get_decoder()

									lowercase_	:   Union[str, Any]  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

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

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

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

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






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






									"""simple docstring"""
									lowercase_	:   Union[str, 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
									lowercase_	:   Union[str, 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 						snake_case__				(							self  )       -> Union[str, Any]:






									"""simple docstring"""
									lowercase_	:   Tuple  				=     self.get_tokenizer()
									# add token to trigger raise
									tokenizer.add_tokens(["""xx"""]  )
									with self.assertRaisesRegex(UpperCamelCase__,   """include"""  ):
																WavaVecaProcessorWithLM(
																    tokenizer=UpperCamelCase__,   feature_extractor=self.get_feature_extractor(),   decoder=self.get_decoder()  )






		def 						snake_case__				(							self  )       -> Tuple:






									"""simple docstring"""
									lowercase_	:   str  				=     self.get_feature_extractor()
									lowercase_	:   List[str]  				=     self.get_tokenizer()
									lowercase_	:   Union[str, Any]  				=     self.get_decoder()

									lowercase_	:   List[Any]  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

									lowercase_	:   List[str]  				=     floats_list((3, 10_00)  )

									lowercase_	:   Any  				=     feature_extractor(UpperCamelCase__,   return_tensors="""np"""  )
									lowercase_	:   str  				=     processor(UpperCamelCase__,   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 						snake_case__				(							self  )       -> Tuple:






									"""simple docstring"""
									lowercase_	:   List[str]  				=     self.get_feature_extractor()
									lowercase_	:   int  				=     self.get_tokenizer()
									lowercase_	:   List[str]  				=     self.get_decoder()

									lowercase_	:   Optional[Any]  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

									lowercase_	:   int  				=     '''This is a test string'''

									lowercase_	:   List[Any]  				=     processor(text=UpperCamelCase__  )

									lowercase_	:   Optional[Any]  				=     tokenizer(UpperCamelCase__  )

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






		def 						snake_case__				(							self,   snake_case__=(2, 10, 16),   snake_case__=77  )       -> List[str]:






									"""simple docstring"""
									np.random.seed(UpperCamelCase__  )
									return np.random.rand(*UpperCamelCase__  )






		def 						snake_case__				(							self  )       -> Tuple:






									"""simple docstring"""
									lowercase_	:   str  				=     self.get_feature_extractor()
									lowercase_	:   Tuple  				=     self.get_tokenizer()
									lowercase_	:   List[Any]  				=     self.get_decoder()

									lowercase_	:   List[str]  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

									lowercase_	:   List[Any]  				=     self._get_dummy_logits(shape=(10, 16),   seed=13  )

									lowercase_	:   List[Any]  				=     processor.decode(UpperCamelCase__  )

									lowercase_	:   Union[str, Any]  				=     decoder.decode_beams(UpperCamelCase__  )[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 						snake_case__				(							self,   snake_case__  )       -> int:






									"""simple docstring"""
									lowercase_	:   List[Any]  				=     self.get_feature_extractor()
									lowercase_	:   Optional[Any]  				=     self.get_tokenizer()
									lowercase_	:   str  				=     self.get_decoder()

									lowercase_	:   Tuple  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

									lowercase_	:   Any  				=     self._get_dummy_logits()

									# note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
									#       otherwise, the LM won't be available to the pool's sub-processes.
									# manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
									if pool_context is None:
																lowercase_	:   Optional[Any]  				=     processor.batch_decode(UpperCamelCase__  )
									else:
																with get_context(UpperCamelCase__  ).Pool() as pool:
																							lowercase_	:   List[str]  				=     processor.batch_decode(UpperCamelCase__,   UpperCamelCase__  )

									lowercase_	:   int  				=     list(UpperCamelCase__  )

									with get_context("""fork"""  ).Pool() as p:
																lowercase_	:   Optional[Any]  				=     decoder.decode_beams_batch(UpperCamelCase__,   UpperCamelCase__  )

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






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






									"""simple docstring"""
									lowercase_	:   Tuple  				=     self.get_feature_extractor()
									lowercase_	:   Any  				=     self.get_tokenizer()
									lowercase_	:   Tuple  				=     self.get_decoder()

									lowercase_	:   str  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

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

									lowercase_	:   Optional[int]  				=     15
									lowercase_	:   Tuple  				=     -20.0
									lowercase_	:   Tuple  				=     -4.0

									lowercase_	:   Tuple  				=     processor.batch_decode(
									    UpperCamelCase__,   beam_width=UpperCamelCase__,   beam_prune_logp=UpperCamelCase__,   token_min_logp=UpperCamelCase__,   )
									lowercase_	:   Optional[Any]  				=     decoded_processor_out.text

									lowercase_	:   List[str]  				=     list(UpperCamelCase__  )

									with get_context("""fork"""  ).Pool() as pool:
																lowercase_	:   Optional[Any]  				=     decoder.decode_beams_batch(
																    UpperCamelCase__,   UpperCamelCase__,   beam_width=UpperCamelCase__,   beam_prune_logp=UpperCamelCase__,   token_min_logp=UpperCamelCase__,   )

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

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

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

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






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     self.get_feature_extractor()
									lowercase_	:   Optional[Any]  				=     self.get_tokenizer()
									lowercase_	:   Any  				=     self.get_decoder()

									lowercase_	:   Optional[int]  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

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

									lowercase_	:   Any  				=     2.0
									lowercase_	:   Any  				=     5.0
									lowercase_	:   Optional[Any]  				=     -20.0
									lowercase_	:   Optional[Any]  				=     True

									lowercase_	:   int  				=     processor.batch_decode(
									    UpperCamelCase__,   alpha=UpperCamelCase__,   beta=UpperCamelCase__,   unk_score_offset=UpperCamelCase__,   lm_score_boundary=UpperCamelCase__,   )
									lowercase_	:   Optional[Any]  				=     decoded_processor_out.text

									lowercase_	:   List[str]  				=     list(UpperCamelCase__  )
									decoder.reset_params(
									    alpha=UpperCamelCase__,   beta=UpperCamelCase__,   unk_score_offset=UpperCamelCase__,   lm_score_boundary=UpperCamelCase__,   )

									with get_context("""fork"""  ).Pool() as pool:
																lowercase_	:   Any  				=     decoder.decode_beams_batch(
																    UpperCamelCase__,   UpperCamelCase__,   )

									lowercase_	:   Optional[int]  				=     [d[0][0] for d in decoded_decoder_out]

									self.assertListEqual(UpperCamelCase__,   UpperCamelCase__  )
									self.assertListEqual(["""<s> </s> <s> </s> </s>""", """</s> </s> <s> </s> </s>"""],   UpperCamelCase__  )
									lowercase_	:   int  				=     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,   UpperCamelCase__  )






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






									"""simple docstring"""
									lowercase_	:   str  				=     WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )

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

									lowercase_	:   str  				=     os.listdir(UpperCamelCase__  )
									lowercase_	:   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(UpperCamelCase__,   UpperCamelCase__  )






		def 						snake_case__				(							self  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     snapshot_download("""hf-internal-testing/processor_with_lm"""  )

									lowercase_	:   int  				=     WavaVecaProcessorWithLM.from_pretrained(UpperCamelCase__  )

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

									lowercase_	:   Tuple  				=     os.listdir(UpperCamelCase__  )
									lowercase_	:   int  				=     os.listdir(UpperCamelCase__  )

									local_decoder_files.sort()
									expected_decoder_files.sort()

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






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






									"""simple docstring"""
									lowercase_	:   List[Any]  				=     WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )
									lowercase_	:   List[Any]  				=     AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""  )

									lowercase_	:   Dict  				=     floats_list((3, 10_00)  )

									lowercase_	:   Optional[int]  				=     processor_wavaveca(UpperCamelCase__,   return_tensors="""np"""  )
									lowercase_	:   Optional[Any]  				=     processor_auto(UpperCamelCase__,   return_tensors="""np"""  )

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

									lowercase_	:   int  				=     self._get_dummy_logits()

									lowercase_	:   List[Any]  				=     processor_wavaveca.batch_decode(UpperCamelCase__  )
									lowercase_	:   List[Any]  				=     processor_auto.batch_decode(UpperCamelCase__  )

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






		def 						snake_case__				(							self  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   str  				=     self.get_feature_extractor()
									lowercase_	:   Optional[int]  				=     self.get_tokenizer()
									lowercase_	:   Any  				=     self.get_decoder()

									lowercase_	:   Union[str, Any]  				=     WavaVecaProcessorWithLM(tokenizer=UpperCamelCase__,   feature_extractor=UpperCamelCase__,   decoder=UpperCamelCase__  )

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






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     [d[key] for d in offsets]
									return retrieved_list






		def 						snake_case__				(							self  )       -> Tuple:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )
									lowercase_	:   Union[str, Any]  				=     self._get_dummy_logits()[0]

									lowercase_	:   Optional[Any]  				=     processor.decode(UpperCamelCase__,   output_word_offsets=UpperCamelCase__  )
									# 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(UpperCamelCase__,   UpperCamelCase__  )  )

									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 						snake_case__				(							self  )       -> str:






									"""simple docstring"""
									lowercase_	:   List[Any]  				=     WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )
									lowercase_	:   Optional[Any]  				=     self._get_dummy_logits()

									lowercase_	:   str  				=     processor.batch_decode(UpperCamelCase__,   output_word_offsets=UpperCamelCase__  )

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

									self.assertListEqual(
									    [""" """.join(self.get_from_offsets(UpperCamelCase__,   """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 						snake_case__				(							self  )       -> List[str]:






									"""simple docstring"""
									import torch

									lowercase_	:   Optional[Any]  				=     load_dataset("""common_voice""",   """en""",   split="""train""",   streaming=UpperCamelCase__  )
									lowercase_	:   Union[str, Any]  				=     ds.cast_column("""audio""",   datasets.Audio(sampling_rate=1_60_00  )  )
									lowercase_	:   List[Any]  				=     iter(UpperCamelCase__  )
									lowercase_	:   List[str]  				=     next(UpperCamelCase__  )

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

									with torch.no_grad():
																lowercase_	:   Any  				=     model(UpperCamelCase__  ).logits.cpu().numpy()

									lowercase_	:   Optional[Any]  				=     processor.decode(logits[0],   output_word_offsets=UpperCamelCase__  )

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

									lowercase_	:   Optional[Any]  				=     '''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(UpperCamelCase__,   """word"""  )  ),   UpperCamelCase__  )
									self.assertEqual(""" """.join(self.get_from_offsets(UpperCamelCase__,   """word"""  )  ),   output.text  )

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

									# fmt: off
									lowercase_	:   Optional[Any]  				=     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_	:   Tuple  				=     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(UpperCamelCase__,   UpperCamelCase__,   atol=0.01  )  )
									self.assertTrue(torch.allclose(UpperCamelCase__,   UpperCamelCase__,   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 DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor, ResNetConfig
from transformers.utils import logging


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




def    _UpperCamelCase						(       SCREAMING_SNAKE_CASE__			)  ->							Tuple:


		'''simple docstring'''


		if "resnet-50" in model_name:
				snake_case  :       int     			=						ResNetConfig.from_pretrained('''microsoft/resnet-50'''			)
		elif "resnet-101" in model_name:
				snake_case  :       Any     			=						ResNetConfig.from_pretrained('''microsoft/resnet-101'''			)
		else:
				raise ValueError('''Model name should include either resnet50 or resnet101'''			)

		snake_case  :       Optional[Any]     			=						DetrConfig(use_timm_backbone=SCREAMING_SNAKE_CASE__ ,       backbone_config=SCREAMING_SNAKE_CASE__			)

		# set label attributes
		snake_case  :       Tuple     			=						'''panoptic''' in model_name
		if is_panoptic:
				snake_case  :       Tuple     			=						250
		else:
				snake_case  :       Any     			=						91
				snake_case  :       str     			=						'''huggingface/label-files'''
				snake_case  :       Optional[Any]     			=						'''coco-detection-id2label.json'''
				snake_case  :       List[str]     			=						json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__ ,       repo_type='''dataset'''			) ,       '''r'''			)			)
				snake_case  :       Optional[Any]     			=						{int(SCREAMING_SNAKE_CASE__			): v for k, v in idalabel.items()}
				snake_case  :       str     			=						idalabel
				snake_case  :       List[Any]     			=						{v: k for k, v in idalabel.items()}

		return config, is_panoptic




def    _UpperCamelCase						(       SCREAMING_SNAKE_CASE__			)  ->							Optional[Any]:


		'''simple docstring'''


		snake_case  :       str     			=						[]

		# stem
		# fmt: off
		rename_keys.append(('''backbone.0.body.conv1.weight''', '''backbone.conv_encoder.model.embedder.embedder.convolution.weight''')			)
		rename_keys.append(('''backbone.0.body.bn1.weight''', '''backbone.conv_encoder.model.embedder.embedder.normalization.weight''')			)
		rename_keys.append(('''backbone.0.body.bn1.bias''', '''backbone.conv_encoder.model.embedder.embedder.normalization.bias''')			)
		rename_keys.append(('''backbone.0.body.bn1.running_mean''', '''backbone.conv_encoder.model.embedder.embedder.normalization.running_mean''')			)
		rename_keys.append(('''backbone.0.body.bn1.running_var''', '''backbone.conv_encoder.model.embedder.embedder.normalization.running_var''')			)
		# stages
		for stage_idx in range(len(config.backbone_config.depths			)			):
				for layer_idx in range(config.backbone_config.depths[stage_idx]			):
						# shortcut
						if layer_idx == 0:
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var',
								    )			)
						# 3 convs
						for i in range(3			):
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean',
								    )			)
								rename_keys.append(
								    (
								        F'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var',
								        F'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var',
								    )			)
    # fmt: on

		for i in range(config.encoder_layers			):
				# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
				rename_keys.append(
				    (
				        F'transformer.encoder.layers.{i}.self_attn.out_proj.weight',
				        F'encoder.layers.{i}.self_attn.out_proj.weight',
				    )			)
				rename_keys.append(
				    (F'transformer.encoder.layers.{i}.self_attn.out_proj.bias', F'encoder.layers.{i}.self_attn.out_proj.bias')			)
				rename_keys.append((F'transformer.encoder.layers.{i}.linear1.weight', F'encoder.layers.{i}.fc1.weight')			)
				rename_keys.append((F'transformer.encoder.layers.{i}.linear1.bias', F'encoder.layers.{i}.fc1.bias')			)
				rename_keys.append((F'transformer.encoder.layers.{i}.linear2.weight', F'encoder.layers.{i}.fc2.weight')			)
				rename_keys.append((F'transformer.encoder.layers.{i}.linear2.bias', F'encoder.layers.{i}.fc2.bias')			)
				rename_keys.append(
				    (F'transformer.encoder.layers.{i}.norm1.weight', F'encoder.layers.{i}.self_attn_layer_norm.weight')			)
				rename_keys.append(
				    (F'transformer.encoder.layers.{i}.norm1.bias', F'encoder.layers.{i}.self_attn_layer_norm.bias')			)
				rename_keys.append(
				    (F'transformer.encoder.layers.{i}.norm2.weight', F'encoder.layers.{i}.final_layer_norm.weight')			)
				rename_keys.append((F'transformer.encoder.layers.{i}.norm2.bias', F'encoder.layers.{i}.final_layer_norm.bias')			)
				# decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
				rename_keys.append(
				    (
				        F'transformer.decoder.layers.{i}.self_attn.out_proj.weight',
				        F'decoder.layers.{i}.self_attn.out_proj.weight',
				    )			)
				rename_keys.append(
				    (F'transformer.decoder.layers.{i}.self_attn.out_proj.bias', F'decoder.layers.{i}.self_attn.out_proj.bias')			)
				rename_keys.append(
				    (
				        F'transformer.decoder.layers.{i}.multihead_attn.out_proj.weight',
				        F'decoder.layers.{i}.encoder_attn.out_proj.weight',
				    )			)
				rename_keys.append(
				    (
				        F'transformer.decoder.layers.{i}.multihead_attn.out_proj.bias',
				        F'decoder.layers.{i}.encoder_attn.out_proj.bias',
				    )			)
				rename_keys.append((F'transformer.decoder.layers.{i}.linear1.weight', F'decoder.layers.{i}.fc1.weight')			)
				rename_keys.append((F'transformer.decoder.layers.{i}.linear1.bias', F'decoder.layers.{i}.fc1.bias')			)
				rename_keys.append((F'transformer.decoder.layers.{i}.linear2.weight', F'decoder.layers.{i}.fc2.weight')			)
				rename_keys.append((F'transformer.decoder.layers.{i}.linear2.bias', F'decoder.layers.{i}.fc2.bias')			)
				rename_keys.append(
				    (F'transformer.decoder.layers.{i}.norm1.weight', F'decoder.layers.{i}.self_attn_layer_norm.weight')			)
				rename_keys.append(
				    (F'transformer.decoder.layers.{i}.norm1.bias', F'decoder.layers.{i}.self_attn_layer_norm.bias')			)
				rename_keys.append(
				    (F'transformer.decoder.layers.{i}.norm2.weight', F'decoder.layers.{i}.encoder_attn_layer_norm.weight')			)
				rename_keys.append(
				    (F'transformer.decoder.layers.{i}.norm2.bias', F'decoder.layers.{i}.encoder_attn_layer_norm.bias')			)
				rename_keys.append(
				    (F'transformer.decoder.layers.{i}.norm3.weight', F'decoder.layers.{i}.final_layer_norm.weight')			)
				rename_keys.append((F'transformer.decoder.layers.{i}.norm3.bias', F'decoder.layers.{i}.final_layer_norm.bias')			)

		# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
		rename_keys.extend(
		    [
		        ('''input_proj.weight''', '''input_projection.weight'''),
		        ('''input_proj.bias''', '''input_projection.bias'''),
		        ('''query_embed.weight''', '''query_position_embeddings.weight'''),
		        ('''transformer.decoder.norm.weight''', '''decoder.layernorm.weight'''),
		        ('''transformer.decoder.norm.bias''', '''decoder.layernorm.bias'''),
		        ('''class_embed.weight''', '''class_labels_classifier.weight'''),
		        ('''class_embed.bias''', '''class_labels_classifier.bias'''),
		        ('''bbox_embed.layers.0.weight''', '''bbox_predictor.layers.0.weight'''),
		        ('''bbox_embed.layers.0.bias''', '''bbox_predictor.layers.0.bias'''),
		        ('''bbox_embed.layers.1.weight''', '''bbox_predictor.layers.1.weight'''),
		        ('''bbox_embed.layers.1.bias''', '''bbox_predictor.layers.1.bias'''),
		        ('''bbox_embed.layers.2.weight''', '''bbox_predictor.layers.2.weight'''),
		        ('''bbox_embed.layers.2.bias''', '''bbox_predictor.layers.2.bias'''),
		    ]			)

		return rename_keys




def    _UpperCamelCase						(       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__			)  ->							List[str]:


		'''simple docstring'''


		snake_case  :       List[str]     			=						state_dict.pop(SCREAMING_SNAKE_CASE__			)
		snake_case  :       List[Any]     			=						val




def    _UpperCamelCase						(       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__=False			)  ->							List[Any]:


		'''simple docstring'''


		snake_case  :       Union[str, Any]     			=						''''''
		if is_panoptic:
				snake_case  :       Optional[Any]     			=						'''detr.'''

		# first: transformer encoder
		for i in range(6			):
				# read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
				snake_case  :       Dict     			=						state_dict.pop(F'{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight'			)
				snake_case  :       Optional[Any]     			=						state_dict.pop(F'{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias'			)
				# next, add query, keys and values (in that order) to the state dict
				snake_case  :       Union[str, Any]     			=						in_proj_weight[:256, :]
				snake_case  :       Any     			=						in_proj_bias[:256]
				snake_case  :       Dict     			=						in_proj_weight[256:512, :]
				snake_case  :       Tuple     			=						in_proj_bias[256:512]
				snake_case  :       Optional[int]     			=						in_proj_weight[-256:, :]
				snake_case  :       Optional[int]     			=						in_proj_bias[-256:]
		# next: transformer decoder (which is a bit more complex because it also includes cross-attention)
		for i in range(6			):
				# read in weights + bias of input projection layer of self-attention
				snake_case  :       Union[str, Any]     			=						state_dict.pop(F'{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight'			)
				snake_case  :       List[str]     			=						state_dict.pop(F'{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias'			)
				# next, add query, keys and values (in that order) to the state dict
				snake_case  :       Union[str, Any]     			=						in_proj_weight[:256, :]
				snake_case  :       str     			=						in_proj_bias[:256]
				snake_case  :       Optional[int]     			=						in_proj_weight[256:512, :]
				snake_case  :       Optional[int]     			=						in_proj_bias[256:512]
				snake_case  :       Optional[Any]     			=						in_proj_weight[-256:, :]
				snake_case  :       Tuple     			=						in_proj_bias[-256:]
				# read in weights + bias of input projection layer of cross-attention
				snake_case  :       Union[str, Any]     			=						state_dict.pop(
				    F'{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight'			)
				snake_case  :       str     			=						state_dict.pop(F'{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias'			)
				# next, add query, keys and values (in that order) of cross-attention to the state dict
				snake_case  :       Optional[Any]     			=						in_proj_weight_cross_attn[:256, :]
				snake_case  :       List[Any]     			=						in_proj_bias_cross_attn[:256]
				snake_case  :       Union[str, Any]     			=						in_proj_weight_cross_attn[256:512, :]
				snake_case  :       Any     			=						in_proj_bias_cross_attn[256:512]
				snake_case  :       List[str]     			=						in_proj_weight_cross_attn[-256:, :]
				snake_case  :       Tuple     			=						in_proj_bias_cross_attn[-256:]




def    _UpperCamelCase						(       )  ->							Dict:


		'''simple docstring'''


		snake_case  :       Optional[Any]     			=						'''http://images.cocodataset.org/val2017/000000039769.jpg'''
		snake_case  :       Any     			=						Image.open(requests.get(SCREAMING_SNAKE_CASE__ ,       stream=SCREAMING_SNAKE_CASE__			).raw			)

		return im




@torch.no_grad()
def    _UpperCamelCase						(       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__=None ,       SCREAMING_SNAKE_CASE__=False			)  ->							Dict:


		'''simple docstring'''


		snake_case			,snake_case  :       Tuple     			=						get_detr_config(SCREAMING_SNAKE_CASE__			)

		# load original model from torch hub
		snake_case  :       Union[str, Any]     			=						{
		    '''detr-resnet-50''': '''detr_resnet50''',
		    '''detr-resnet-101''': '''detr_resnet101''',
		}
		logger.info(F'Converting model {model_name}...'			)
		snake_case  :       Union[str, Any]     			=						torch.hub.load('''facebookresearch/detr''' ,       model_name_to_original_name[model_name] ,       pretrained=SCREAMING_SNAKE_CASE__			).eval()
		snake_case  :       Optional[int]     			=						detr.state_dict()
		# rename keys
		for src, dest in create_rename_keys(SCREAMING_SNAKE_CASE__			):
				if is_panoptic:
						snake_case  :       Union[str, Any]     			=						'''detr.''' + src
				rename_key(SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__			)
		# query, key and value matrices need special treatment
		read_in_q_k_v(SCREAMING_SNAKE_CASE__ ,       is_panoptic=SCREAMING_SNAKE_CASE__			)
		# important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
		snake_case  :       List[str]     			=						'''detr.model.''' if is_panoptic else '''model.'''
		for key in state_dict.copy().keys():
				if is_panoptic:
						if (
						    key.startswith('''detr'''			)
						    and not key.startswith('''class_labels_classifier'''			)
						    and not key.startswith('''bbox_predictor'''			)
						):
								snake_case  :       Optional[Any]     			=						state_dict.pop(SCREAMING_SNAKE_CASE__			)
								snake_case  :       Optional[int]     			=						val
						elif "class_labels_classifier" in key or "bbox_predictor" in key:
								snake_case  :       Dict     			=						state_dict.pop(SCREAMING_SNAKE_CASE__			)
								snake_case  :       List[Any]     			=						val
						elif key.startswith('''bbox_attention'''			) or key.startswith('''mask_head'''			):
								continue
						else:
								snake_case  :       int     			=						state_dict.pop(SCREAMING_SNAKE_CASE__			)
								snake_case  :       Optional[int]     			=						val
				else:
						if not key.startswith('''class_labels_classifier'''			) and not key.startswith('''bbox_predictor'''			):
								snake_case  :       Dict     			=						state_dict.pop(SCREAMING_SNAKE_CASE__			)
								snake_case  :       Any     			=						val

    # finally, create HuggingFace model and load state dict
		snake_case  :       Optional[Any]     			=						DetrForSegmentation(SCREAMING_SNAKE_CASE__			) if is_panoptic else DetrForObjectDetection(SCREAMING_SNAKE_CASE__			)
		model.load_state_dict(SCREAMING_SNAKE_CASE__			)
		model.eval()

		# verify our conversion on an image
		snake_case  :       List[str]     			=						'''coco_panoptic''' if is_panoptic else '''coco_detection'''
		snake_case  :       Optional[int]     			=						DetrImageProcessor(format=SCREAMING_SNAKE_CASE__			)

		snake_case  :       int     			=						processor(images=prepare_img() ,       return_tensors='''pt'''			)
		snake_case  :       Optional[int]     			=						encoding['''pixel_values''']

		snake_case  :       str     			=						detr(SCREAMING_SNAKE_CASE__			)
		snake_case  :       Optional[int]     			=						model(SCREAMING_SNAKE_CASE__			)

		assert torch.allclose(outputs.logits ,       original_outputs['''pred_logits'''] ,       atol=1E-3			)
		assert torch.allclose(outputs.pred_boxes ,       original_outputs['''pred_boxes'''] ,       atol=1E-3			)
		if is_panoptic:
				assert torch.allclose(outputs.pred_masks ,       original_outputs['''pred_masks'''] ,       atol=1E-4			)
		print('''Looks ok!'''			)

		if pytorch_dump_folder_path is not None:
				# Save model and image processor
				logger.info(F'Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'			)
				Path(SCREAMING_SNAKE_CASE__			).mkdir(exist_ok=SCREAMING_SNAKE_CASE__			)
				model.save_pretrained(SCREAMING_SNAKE_CASE__			)
				processor.save_pretrained(SCREAMING_SNAKE_CASE__			)

		if push_to_hub:
				# Upload model and image processor to the hub
				logger.info('''Uploading PyTorch model and image processor to the hub...'''			)
				model.push_to_hub(F'nielsr/{model_name}'			)
				processor.push_to_hub(F'nielsr/{model_name}'			)


if __name__ == "__main__":
						lowercase__						       =      argparse.ArgumentParser()

						parser.add_argument(
						    "--model_name",
						    default="detr-resnet-50",
						    type=str,
						    choices=["detr-resnet-50", "detr-resnet-101"],
						    help="Name of the DETR model you'd like to convert.",
						)
						parser.add_argument(
						    "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
						)
						parser.add_argument("--push_to_hub", action="store_true", help="Whether to push the model to the hub or not.")
						lowercase__						       =      parser.parse_args()
						convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)


