





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




              """simple docstring"""


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

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

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

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

              # load decoder from hub
              UpperCamelCase_          =  "hf-internal-testing/ngram-beam-search-decoder"

       def lowercase (		self: str					,   **_SCREAMING_SNAKE_CASE: Any		)       -> List[str]:




              """simple docstring"""


              UpperCamelCase_          =  self.add_kwargs_tokens_map.copy()
              kwargs.update(_SCREAMING_SNAKE_CASE		)
              return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname					,   **_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Optional[int]					,   **_SCREAMING_SNAKE_CASE: Optional[int]		)       -> str:




              """simple docstring"""


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

       def lowercase (		self: Optional[int]					,   **_SCREAMING_SNAKE_CASE: List[Any]		)       -> List[str]:




              """simple docstring"""


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

       def lowercase (		self: int		)       -> Optional[Any]:




              """simple docstring"""


              shutil.rmtree(self.tmpdirname		)

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




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

              processor.save_pretrained(self.tmpdirname		)
              UpperCamelCase_          =  WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname		)

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

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

              # decoder
              self.assertEqual(processor.decoder._alphabet.labels					,   decoder._alphabet.labels		)
              self.assertEqual(
                  processor.decoder.model_container[decoder._model_key]._unigram_set					,   decoder.model_container[decoder._model_key]._unigram_set					,   )
              self.assertIsInstance(processor.decoder					,   _SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Optional[Any]		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  WavaVecaProcessorWithLM(
                  tokenizer=self.get_tokenizer()					,   feature_extractor=self.get_feature_extractor()					,   decoder=self.get_decoder()		)
              processor.save_pretrained(self.tmpdirname		)

              # make sure that error is thrown when decoder alphabet doesn't match
              UpperCamelCase_          =  WavaVecaProcessorWithLM.from_pretrained(
                  self.tmpdirname					,   alpha=5.0					,   beta=3.0					,   score_boundary=-7.0					,   unk_score_offset=3		)

              # decoder
              self.assertEqual(processor.language_model.alpha					,   5.0		)
              self.assertEqual(processor.language_model.beta					,   3.0		)
              self.assertEqual(processor.language_model.score_boundary					,   -7.0		)
              self.assertEqual(processor.language_model.unk_score_offset					,   3		)

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




              """simple docstring"""


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

       def lowercase (		self: int		)       -> Tuple:




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  floats_list((3, 1000)		)

              UpperCamelCase_          =  feature_extractor(_SCREAMING_SNAKE_CASE					,   return_tensors="np"		)
              UpperCamelCase_          =  processor(_SCREAMING_SNAKE_CASE					,   return_tensors="np"		)

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

       def lowercase (		self: Any		)       -> Tuple:




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  "This is a test string"

              UpperCamelCase_          =  processor(text=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  tokenizer(_SCREAMING_SNAKE_CASE		)

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

       def lowercase (		self: Any					,   _SCREAMING_SNAKE_CASE: List[str]=(2, 10, 16)					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=77		)       -> Dict:




              """simple docstring"""


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

       def lowercase (		self: Optional[Any]		)       -> int:




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

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

              UpperCamelCase_          =  processor.decode(_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  decoder.decode_beams(_SCREAMING_SNAKE_CASE		)[0]

              self.assertEqual(decoded_decoder[0]					,   decoded_processor.text		)
              self.assertEqual("</s> <s> </s>"					,   decoded_processor.text		)
              self.assertEqual(decoded_decoder[-2]					,   decoded_processor.logit_score		)
              self.assertEqual(decoded_decoder[-1]					,   decoded_processor.lm_score		)

       @parameterized.expand([[None], ["fork"], ["spawn"]]		)
       def lowercase (		self: List[str]					,   _SCREAMING_SNAKE_CASE: Tuple		)       -> int:




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  self._get_dummy_logits()

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

              UpperCamelCase_          =  list(_SCREAMING_SNAKE_CASE		)

              with get_context("fork"		).Pool() as p:
                     UpperCamelCase_          =  decoder.decode_beams_batch(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

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

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

       def lowercase (		self: List[Any]		)       -> Tuple:




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  self._get_dummy_logits()

              UpperCamelCase_          =  15
              UpperCamelCase_          =  -20.0
              UpperCamelCase_          =  -4.0

              UpperCamelCase_          =  processor.batch_decode(
                  _SCREAMING_SNAKE_CASE					,   beam_width=_SCREAMING_SNAKE_CASE					,   beam_prune_logp=_SCREAMING_SNAKE_CASE					,   token_min_logp=_SCREAMING_SNAKE_CASE					,   )
              UpperCamelCase_          =  decoded_processor_out.text

              UpperCamelCase_          =  list(_SCREAMING_SNAKE_CASE		)

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

              UpperCamelCase_          =  [d[0][0] for d in decoded_decoder_out]
              UpperCamelCase_          =  [d[0][2] for d in decoded_decoder_out]
              UpperCamelCase_          =  [d[0][3] for d in decoded_decoder_out]

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

              self.assertTrue(np.array_equal(_SCREAMING_SNAKE_CASE					,   decoded_processor_out.logit_score		)		)
              self.assertTrue(np.allclose([-20.0_54, -18.4_47]					,   _SCREAMING_SNAKE_CASE					,   atol=1e-3		)		)

              self.assertTrue(np.array_equal(_SCREAMING_SNAKE_CASE					,   decoded_processor_out.lm_score		)		)
              self.assertTrue(np.allclose([-15.5_54, -13.94_74]					,   _SCREAMING_SNAKE_CASE					,   atol=1e-3		)		)

       def lowercase (		self: int		)       -> List[Any]:




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  self._get_dummy_logits()

              UpperCamelCase_          =  2.0
              UpperCamelCase_          =  5.0
              UpperCamelCase_          =  -20.0
              UpperCamelCase_          =  True

              UpperCamelCase_          =  processor.batch_decode(
                  _SCREAMING_SNAKE_CASE					,   alpha=_SCREAMING_SNAKE_CASE					,   beta=_SCREAMING_SNAKE_CASE					,   unk_score_offset=_SCREAMING_SNAKE_CASE					,   lm_score_boundary=_SCREAMING_SNAKE_CASE					,   )
              UpperCamelCase_          =  decoded_processor_out.text

              UpperCamelCase_          =  list(_SCREAMING_SNAKE_CASE		)
              decoder.reset_params(
                  alpha=_SCREAMING_SNAKE_CASE					,   beta=_SCREAMING_SNAKE_CASE					,   unk_score_offset=_SCREAMING_SNAKE_CASE					,   lm_score_boundary=_SCREAMING_SNAKE_CASE					,   )

              with get_context("fork"		).Pool() as pool:
                     UpperCamelCase_          =  decoder.decode_beams_batch(
                         _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   )

              UpperCamelCase_          =  [d[0][0] for d in decoded_decoder_out]

              self.assertListEqual(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)
              self.assertListEqual(["<s> </s> <s> </s> </s>", "</s> </s> <s> </s> </s>"]					,   _SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  processor.decoder.model_container[processor.decoder._model_key]
              self.assertEqual(lm_model.alpha					,   2.0		)
              self.assertEqual(lm_model.beta					,   5.0		)
              self.assertEqual(lm_model.unk_score_offset					,   -20.0		)
              self.assertEqual(lm_model.score_boundary					,   _SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Dict		)       -> Tuple:




              """simple docstring"""


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

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

              UpperCamelCase_          =  os.listdir(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  ["alphabet.json", "language_model"]

              downloaded_decoder_files.sort()
              expected_decoder_files.sort()

              # test that only decoder relevant files from
              # https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
              # are downloaded and none of the rest (e.g. README.md, ...)
              self.assertListEqual(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

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




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM.from_pretrained(_SCREAMING_SNAKE_CASE		)

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

              UpperCamelCase_          =  os.listdir(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  os.listdir(_SCREAMING_SNAKE_CASE		)

              local_decoder_files.sort()
              expected_decoder_files.sort()

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

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




              """simple docstring"""


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

              UpperCamelCase_          =  floats_list((3, 1000)		)

              UpperCamelCase_          =  processor_wavaveca(_SCREAMING_SNAKE_CASE					,   return_tensors="np"		)
              UpperCamelCase_          =  processor_auto(_SCREAMING_SNAKE_CASE					,   return_tensors="np"		)

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

              UpperCamelCase_          =  self._get_dummy_logits()

              UpperCamelCase_          =  processor_wavaveca.batch_decode(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  processor_auto.batch_decode(_SCREAMING_SNAKE_CASE		)

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

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




              """simple docstring"""


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

              UpperCamelCase_          =  WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE					,   feature_extractor=_SCREAMING_SNAKE_CASE					,   decoder=_SCREAMING_SNAKE_CASE		)

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

       @staticmethod
       def lowercase (		_SCREAMING_SNAKE_CASE: Union[str, Any]					,   _SCREAMING_SNAKE_CASE: List[str]		)       -> Any:




              """simple docstring"""


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

       def lowercase (		self: int		)       -> List[Any]:




              """simple docstring"""


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

              UpperCamelCase_          =  processor.decode(_SCREAMING_SNAKE_CASE					,   output_word_offsets=_SCREAMING_SNAKE_CASE		)
              # check Wav2Vec2CTCTokenizerOutput keys for word
              self.assertEqual(len(outputs.keys()		)					,   4		)
              self.assertTrue("text" in outputs		)
              self.assertTrue("word_offsets" in outputs		)
              self.assertTrue(isinstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)		)

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

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




              """simple docstring"""


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

              UpperCamelCase_          =  processor.batch_decode(_SCREAMING_SNAKE_CASE					,   output_word_offsets=_SCREAMING_SNAKE_CASE		)

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

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



       @slow
       @require_torch
       @require_torchaudio
       def lowercase (		self: int		)       -> Union[str, Any]:




              """simple docstring"""


              import torch

              UpperCamelCase_          =  load_dataset("common_voice"					,   "en"					,   split="train"					,   streaming=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  ds.cast_column("audio"					,   datasets.Audio(sampling_rate=16000		)		)
              UpperCamelCase_          =  iter(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  next(_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  AutoProcessor.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm"		)
              UpperCamelCase_          =  WavaVecaForCTC.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm"		)

              # compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
              UpperCamelCase_          =  processor(sample["audio"]["array"]					,   return_tensors="pt"		).input_values

              with torch.no_grad():
                     UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE		).logits.cpu().numpy()

              UpperCamelCase_          =  processor.decode(logits[0]					,   output_word_offsets=_SCREAMING_SNAKE_CASE		)

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

              UpperCamelCase_          =  "WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL"

              # output words
              self.assertEqual(" ".join(self.get_from_offsets(_SCREAMING_SNAKE_CASE					,   "word"		)		)					,   _SCREAMING_SNAKE_CASE		)
              self.assertEqual(" ".join(self.get_from_offsets(_SCREAMING_SNAKE_CASE					,   "word"		)		)					,   output.text		)

              # output times
              UpperCamelCase_          =  torch.tensor(self.get_from_offsets(_SCREAMING_SNAKE_CASE					,   "start_time"		)		)
              UpperCamelCase_          =  torch.tensor(self.get_from_offsets(_SCREAMING_SNAKE_CASE					,   "end_time"		)		)

              # fmt: off
              UpperCamelCase_          =  torch.tensor([1.41_99, 1.65_99, 2.25_99, 3.0, 3.24, 3.59_99, 3.79_99, 4.09_99, 4.26, 4.94, 5.28, 5.65_99, 5.78, 5.94, 6.32, 6.53_99, 6.65_99]		)
              UpperCamelCase_          =  torch.tensor([1.53_99, 1.89_99, 2.9, 3.16, 3.53_99, 3.72, 4.01_99, 4.17_99, 4.76, 5.15_99, 5.55_99, 5.69_99, 5.86, 6.19_99, 6.38, 6.61_99, 6.94]		)
              # fmt: on

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








# Copyright (c) 2021-, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

####################################################################################################

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_=0  )     ->		int:
       # Format the message.
       if name is None:
              UpperCamelCase_          =  None
       else:
              UpperCamelCase_          =  "." * max(0					,		spaces - 2  ) + "# {:" + str(50 - spaces  ) + "s}"
              UpperCamelCase_          =  fmt.format(UpperCamelCase_  )

       # Print and recurse (if needed).
       if isinstance(UpperCamelCase_					,		UpperCamelCase_  ):
              if msg is not None:
                     print(UpperCamelCase_  )
              for k in val.keys():
                     recursive_print(UpperCamelCase_					,		val[k]					,		spaces + 2  )
       elif isinstance(UpperCamelCase_					,		torch.Tensor  ):
              print(UpperCamelCase_					,		":"					,		val.size()  )
       else:
              print(UpperCamelCase_					,		":"					,		UpperCamelCase_  )




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )     ->		int:
       # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
       # for compatibility with later versions of NVIDIA Megatron-LM.
       # The inverse operation is performed inside Megatron-LM to read checkpoints:
       # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
       # If param is the weight tensor of the self-attention block, the returned tensor
       # will have to be transposed one more time to be read by HuggingFace GPT2.
       UpperCamelCase_          =  param.size()
       if checkpoint_version == 1.0:
              # version 1.0 stores [num_heads * hidden_size * num_splits, :]
              UpperCamelCase_          =  (num_heads, hidden_size, num_splits) + input_shape[1:]
              UpperCamelCase_          =  param.view(*UpperCamelCase_  )
              UpperCamelCase_          =  param.transpose(0					,		2  )
              UpperCamelCase_          =  param.transpose(1					,		2  ).contiguous()
       elif checkpoint_version >= 2.0:
              # other versions store [num_heads * num_splits * hidden_size, :]
              UpperCamelCase_          =  (num_heads, num_splits, hidden_size) + input_shape[1:]
              UpperCamelCase_          =  param.view(*UpperCamelCase_  )
              UpperCamelCase_          =  param.transpose(0					,		1  ).contiguous()
       UpperCamelCase_          =  param.view(*UpperCamelCase_  )
       return param




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )     ->		int:
       # The converted output model.
       UpperCamelCase_          =  {}

       # old versions did not store training args
       UpperCamelCase_          =  input_state_dict.get("args"					,		UpperCamelCase_  )
       if ds_args is not None:
              # do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
              # from pprint import pprint
              # pprint(vars(ds_args))

              UpperCamelCase_          =  ds_args.padded_vocab_size
              UpperCamelCase_          =  ds_args.max_position_embeddings
              UpperCamelCase_          =  ds_args.hidden_size
              UpperCamelCase_          =  ds_args.num_layers
              UpperCamelCase_          =  ds_args.num_attention_heads
              UpperCamelCase_          =  ds_args.ffn_hidden_size
              # pprint(config)

       # The number of heads.
       UpperCamelCase_          =  config.n_head
       # The hidden_size per head.
       UpperCamelCase_          =  config.n_embd // config.n_head
       # Megatron-LM checkpoint version
       if "checkpoint_version" in input_state_dict.keys():
              UpperCamelCase_          =  input_state_dict["checkpoint_version"]
       else:
              UpperCamelCase_          =  0.0

       # The model.
       UpperCamelCase_          =  input_state_dict["model"]
       # The language model.
       UpperCamelCase_          =  model["language_model"]
       # The embeddings.
       UpperCamelCase_          =  lm["embedding"]

       # The word embeddings.
       UpperCamelCase_          =  embeddings["word_embeddings"]["weight"]
       # Truncate the embedding table to vocab_size rows.
       UpperCamelCase_          =  word_embeddings[: config.vocab_size, :]
       UpperCamelCase_          =  word_embeddings

       # The position embeddings.
       UpperCamelCase_          =  embeddings["position_embeddings"]["weight"]
       # Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
       UpperCamelCase_          =  pos_embeddings.size(0  )
       if n_positions != config.n_positions:
              raise ValueError(
                  F'''pos_embeddings.max_sequence_length={n_positions} and config.n_positions={config.n_positions} don\'t match'''  )
       # Store the position embeddings.
       UpperCamelCase_          =  pos_embeddings

       # The transformer.
       UpperCamelCase_          =  lm["transformer"] if "transformer" in lm.keys() else lm["encoder"]

       # The regex to extract layer names.
       UpperCamelCase_          =  re.compile(r"layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)"  )

       # The simple map of names for "automated" rules.
       UpperCamelCase_          =  {
           "attention.dense": ".attn.c_proj.",
           "self_attention.dense": ".attn.c_proj.",
           "mlp.dense_h_to_4h": ".mlp.c_fc.",
           "mlp.dense_4h_to_h": ".mlp.c_proj.",
       }

       # Extract the layers.
       for key, val in transformer.items():
              # Match the name.
              UpperCamelCase_          =  layer_re.match(UpperCamelCase_  )

              # Stop if that's not a layer
              if m is None:
                     break

              # The index of the layer.
              UpperCamelCase_          =  int(m.group(1  )  )
              # The name of the operation.
              UpperCamelCase_          =  m.group(2  )
              # Is it a weight or a bias?
              UpperCamelCase_          =  m.group(3  )

              # The name of the layer.
              UpperCamelCase_          =  F'''transformer.h.{layer_idx}'''

              # For layernorm(s), simply store the layer norm.
              if op_name.endswith("layernorm"  ):
                     UpperCamelCase_          =  "ln_1" if op_name.startswith("input"  ) else "ln_2"
                     UpperCamelCase_          =  val

              # Transpose the QKV matrix.
              elif (
                  op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
              ) and weight_or_bias == "weight":
                     # Insert a tensor of 1x1xDxD bias.
                     UpperCamelCase_          =  torch.tril(torch.ones((n_positions, n_positions)					,		dtype=torch.floataa  )  ).view(
                         1					,		1					,		UpperCamelCase_					,		UpperCamelCase_  )
                     UpperCamelCase_          =  causal_mask

                     # Insert a "dummy" tensor for masked_bias.
                     UpperCamelCase_          =  torch.tensor(-1e4					,		dtype=torch.floataa  )
                     UpperCamelCase_          =  masked_bias

                     UpperCamelCase_          =  fix_query_key_value_ordering(UpperCamelCase_					,		UpperCamelCase_					,		3					,		UpperCamelCase_					,		UpperCamelCase_  )
                     # Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
                     UpperCamelCase_          =  out_val.transpose(0					,		1  ).contiguous()
                     # Store.
                     UpperCamelCase_          =  out_val

              # Transpose the bias.
              elif (
                  op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
              ) and weight_or_bias == "bias":
                     UpperCamelCase_          =  fix_query_key_value_ordering(UpperCamelCase_					,		UpperCamelCase_					,		3					,		UpperCamelCase_					,		UpperCamelCase_  )
                     # Store. No change of shape.
                     UpperCamelCase_          =  out_val

              # Transpose the weights.
              elif weight_or_bias == "weight":
                     UpperCamelCase_          =  megatron_to_transformers[op_name]
                     UpperCamelCase_          =  val.transpose(0					,		1  )

              # Copy the bias.
              elif weight_or_bias == "bias":
                     UpperCamelCase_          =  megatron_to_transformers[op_name]
                     UpperCamelCase_          =  val

    # DEBUG.
       assert config.n_layer == layer_idx + 1

       # The final layernorm.
       UpperCamelCase_          =  transformer["final_layernorm.weight"]
       UpperCamelCase_          =  transformer["final_layernorm.bias"]

       # For LM head, transformers' wants the matrix to weight embeddings.
       UpperCamelCase_          =  word_embeddings

       # It should be done!
       return output_state_dict




def  lowerCAmelCase_ (							)     ->		Union[str, Any]:
       # Create the argument parser.
       UpperCamelCase_          =  argparse.ArgumentParser()
       parser.add_argument("--print-checkpoint-structure"					,		action="store_true"  )
       parser.add_argument(
           "path_to_checkpoint"					,		type=UpperCamelCase_					,		help="Path to the checkpoint file (.zip archive or direct .pt file)"					,		)
       parser.add_argument(
           "--config_file"					,		default=""					,		type=UpperCamelCase_					,		help="An optional config json file describing the pre-trained model."					,		)
       UpperCamelCase_          =  parser.parse_args()

       # Extract the basename.
       UpperCamelCase_          =  os.path.dirname(args.path_to_checkpoint  )

       # Load the model.
       # the .zip is very optional, let's keep it for backward compatibility
       print(F'''Extracting PyTorch state dictionary from {args.path_to_checkpoint}'''  )
       if args.path_to_checkpoint.endswith(".zip"  ):
              with zipfile.ZipFile(args.path_to_checkpoint					,		"r"  ) as checkpoint:
                     with checkpoint.open("release/mp_rank_00/model_optim_rng.pt"  ) as pytorch_dict:
                            UpperCamelCase_          =  torch.load(UpperCamelCase_					,		map_location="cpu"  )
       else:
              UpperCamelCase_          =  torch.load(args.path_to_checkpoint					,		map_location="cpu"  )

       UpperCamelCase_          =  input_state_dict.get("args"					,		UpperCamelCase_  )

       # Read the config, or default to the model released by NVIDIA.
       if args.config_file == "":
              if ds_args is not None:
                     if ds_args.bias_gelu_fusion:
                            UpperCamelCase_          =  "gelu_fast"
                     elif ds_args.openai_gelu:
                            UpperCamelCase_          =  "gelu_new"
                     else:
                            UpperCamelCase_          =  "gelu"
              else:
                     # in the very early days this used to be "gelu_new"
                     UpperCamelCase_          =  "gelu_new"

              # Spell out all parameters in case the defaults change.
              UpperCamelCase_          =  GPTaConfig(
                  vocab_size=50257					,		n_positions=1024					,		n_embd=1024					,		n_layer=24					,		n_head=16					,		n_inner=4096					,		activation_function=UpperCamelCase_					,		resid_pdrop=0.1					,		embd_pdrop=0.1					,		attn_pdrop=0.1					,		layer_norm_epsilon=1e-5					,		initializer_range=0.02					,		summary_type="cls_index"					,		summary_use_proj=UpperCamelCase_					,		summary_activation=UpperCamelCase_					,		summary_proj_to_labels=UpperCamelCase_					,		summary_first_dropout=0.1					,		scale_attn_weights=UpperCamelCase_					,		use_cache=UpperCamelCase_					,		bos_token_id=50256					,		eos_token_id=50256					,		)
       else:
              UpperCamelCase_          =  GPTaConfig.from_json_file(args.config_file  )

       UpperCamelCase_          =  ["GPT2LMHeadModel"]

       # Convert.
       print("Converting"  )
       UpperCamelCase_          =  convert_megatron_checkpoint(UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )

       # Print the structure of converted state dict.
       if args.print_checkpoint_structure:
              recursive_print(UpperCamelCase_					,		UpperCamelCase_  )

       # Add tokenizer class info to config
       # see https://github.com/huggingface/transformers/issues/13906)
       if ds_args is not None:
              UpperCamelCase_          =  ds_args.tokenizer_type
              if tokenizer_type == "GPT2BPETokenizer":
                     UpperCamelCase_          =  "gpt2"
              elif tokenizer_type == "PretrainedFromHF":
                     UpperCamelCase_          =  ds_args.tokenizer_name_or_path
              else:
                     raise ValueError(F'''Unrecognized tokenizer_type {tokenizer_type}'''  )
       else:
              UpperCamelCase_          =  "gpt2"

       UpperCamelCase_          =  AutoTokenizer.from_pretrained(UpperCamelCase_  )
       UpperCamelCase_          =  type(UpperCamelCase_  ).__name__
       UpperCamelCase_          =  tokenizer_class

       # Store the config to file.
       print("Saving config"  )
       config.save_pretrained(UpperCamelCase_  )

       # Save tokenizer based on args
       print(F'''Adding {tokenizer_class} tokenizer files'''  )
       tokenizer.save_pretrained(UpperCamelCase_  )

       # Store the state_dict to file.
       UpperCamelCase_          =  os.path.join(UpperCamelCase_					,		"pytorch_model.bin"  )
       print(F'''Saving checkpoint to "{output_checkpoint_file}"'''  )
       torch.save(UpperCamelCase_					,		UpperCamelCase_  )


####################################################################################################

if __name__ == "__main__":
     main()

####################################################################################################

