






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




 def      __magic_name__	(      self	:							int      )	->						List[str]:
       SCREAMING_SNAKE_CASE__     :				Dict									='''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
       SCREAMING_SNAKE_CASE__     :				int									=dict(zip(__lowercase					,      range(len(__lowercase      )      )      )      )

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

       SCREAMING_SNAKE_CASE__     :				Any									=tempfile.mkdtemp()
       SCREAMING_SNAKE_CASE__     :				Tuple									=os.path.join(self.tmpdirname					,      VOCAB_FILES_NAMES['''vocab_file''']      )
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=os.path.join(self.tmpdirname					,      __lowercase      )
       with open(self.vocab_file					,      '''w'''					,      encoding='''utf-8'''      ) as fp:
             fp.write(json.dumps(__lowercase      ) + '''\n'''      )

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

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




 def      __magic_name__	(      self	:							str					,      **__lowercase	:							List[Any]      )	->						List[str]:
       SCREAMING_SNAKE_CASE__     :				List[str]									=self.add_kwargs_tokens_map.copy()
       kwargs.update(__lowercase      )
       return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname					,      **__lowercase      )




 def      __magic_name__	(      self	:							int					,      **__lowercase	:							Dict      )	->						Any:
       return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname					,      **__lowercase      )




 def      __magic_name__	(      self	:							int					,      **__lowercase	:							int      )	->						int:
       return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name					,      **__lowercase      )




 def      __magic_name__	(      self	:							str      )	->						Optional[int]:
       shutil.rmtree(self.tmpdirname      )




 def      __magic_name__	(      self	:							str      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				List[str]									=self.get_tokenizer()
       SCREAMING_SNAKE_CASE__     :				List[str]									=self.get_feature_extractor()
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.get_decoder()

       SCREAMING_SNAKE_CASE__     :				List[str]									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

       processor.save_pretrained(self.tmpdirname      )
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname      )

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

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

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




 def      __magic_name__	(      self	:							Any      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				List[str]									=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
       SCREAMING_SNAKE_CASE__     :				Tuple									=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      __magic_name__	(      self	:							List[Any]      )	->						str:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.get_tokenizer()
       # add token to trigger raise
       tokenizer.add_tokens(['''xx''']      )
       with self.assertRaisesRegex(__lowercase					,      '''include'''      ):
             WavaVecaProcessorWithLM(
                 tokenizer=__lowercase					,      feature_extractor=self.get_feature_extractor()					,      decoder=self.get_decoder()      )




 def      __magic_name__	(      self	:							List[str]      )	->						Optional[Any]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=self.get_feature_extractor()
       SCREAMING_SNAKE_CASE__     :				Dict									=self.get_tokenizer()
       SCREAMING_SNAKE_CASE__     :				str									=self.get_decoder()

       SCREAMING_SNAKE_CASE__     :				Tuple									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

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

       SCREAMING_SNAKE_CASE__     :				Dict									=feature_extractor(__lowercase					,      return_tensors='''np'''      )
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=processor(__lowercase					,      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      __magic_name__	(      self	:							Tuple      )	->						List[str]:
       SCREAMING_SNAKE_CASE__     :				Any									=self.get_feature_extractor()
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.get_tokenizer()
       SCREAMING_SNAKE_CASE__     :				List[str]									=self.get_decoder()

       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

       SCREAMING_SNAKE_CASE__     :				List[str]									='''This is a test string'''

       SCREAMING_SNAKE_CASE__     :				Any									=processor(text=__lowercase      )

       SCREAMING_SNAKE_CASE__     :				str									=tokenizer(__lowercase      )

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




 def      __magic_name__	(      self	:							List[str]					,      __lowercase	:							Tuple=(2, 10, 16)					,      __lowercase	:							Union[str, Any]=77      )	->						str:
       np.random.seed(__lowercase      )
       return np.random.rand(*__lowercase      )




 def      __magic_name__	(      self	:							Optional[int]      )	->						int:
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.get_feature_extractor()
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.get_tokenizer()
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.get_decoder()

       SCREAMING_SNAKE_CASE__     :				str									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

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

       SCREAMING_SNAKE_CASE__     :				int									=processor.decode(__lowercase      )

       SCREAMING_SNAKE_CASE__     :				Any									=decoder.decode_beams(__lowercase      )[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      __magic_name__	(      self	:							int					,      __lowercase	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Dict									=self.get_feature_extractor()
       SCREAMING_SNAKE_CASE__     :				Dict									=self.get_tokenizer()
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.get_decoder()

       SCREAMING_SNAKE_CASE__     :				List[str]									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

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

       SCREAMING_SNAKE_CASE__     :				Dict									=list(__lowercase      )

       with get_context('''fork'''      ).Pool() as p:
             SCREAMING_SNAKE_CASE__     :				Dict									=decoder.decode_beams_batch(__lowercase					,      __lowercase      )

       SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__     :				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(__lowercase					,      decoded_processor.text      )
       self.assertListEqual(['''<s> <s> </s>''', '''<s> <s> <s>''']					,      decoded_processor.text      )
       self.assertListEqual(__lowercase					,      decoded_processor.logit_score      )
       self.assertListEqual(__lowercase					,      decoded_processor.lm_score      )




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

       SCREAMING_SNAKE_CASE__     :				int									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

       SCREAMING_SNAKE_CASE__     :				Any									=self._get_dummy_logits()

       SCREAMING_SNAKE_CASE__     :				int									=15
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=-20.0
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=-4.0

       SCREAMING_SNAKE_CASE__     :				str									=processor.batch_decode(
           __lowercase					,      beam_width=__lowercase					,      beam_prune_logp=__lowercase					,      token_min_logp=__lowercase					,      )
       SCREAMING_SNAKE_CASE__     :				str									=decoded_processor_out.text

       SCREAMING_SNAKE_CASE__     :				Optional[int]									=list(__lowercase      )

       with get_context('''fork'''      ).Pool() as pool:
             SCREAMING_SNAKE_CASE__     :				List[Any]									=decoder.decode_beams_batch(
                 __lowercase					,      __lowercase					,      beam_width=__lowercase					,      beam_prune_logp=__lowercase					,      token_min_logp=__lowercase					,      )

       SCREAMING_SNAKE_CASE__     :				str									=[d[0][0] for d in decoded_decoder_out]
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=[d[0][2] for d in decoded_decoder_out]
       SCREAMING_SNAKE_CASE__     :				Tuple									=[d[0][3] for d in decoded_decoder_out]

       self.assertListEqual(__lowercase					,      __lowercase      )
       self.assertListEqual(['''</s> <s> <s>''', '''<s> <s> <s>''']					,      __lowercase      )

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

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




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

       SCREAMING_SNAKE_CASE__     :				Any									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

       SCREAMING_SNAKE_CASE__     :				int									=self._get_dummy_logits()

       SCREAMING_SNAKE_CASE__     :				Tuple									=2.0
       SCREAMING_SNAKE_CASE__     :				List[str]									=5.0
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=-20.0
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=True

       SCREAMING_SNAKE_CASE__     :				Dict									=processor.batch_decode(
           __lowercase					,      alpha=__lowercase					,      beta=__lowercase					,      unk_score_offset=__lowercase					,      lm_score_boundary=__lowercase					,      )
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=decoded_processor_out.text

       SCREAMING_SNAKE_CASE__     :				int									=list(__lowercase      )
       decoder.reset_params(
           alpha=__lowercase					,      beta=__lowercase					,      unk_score_offset=__lowercase					,      lm_score_boundary=__lowercase					,      )

       with get_context('''fork'''      ).Pool() as pool:
             SCREAMING_SNAKE_CASE__     :				int									=decoder.decode_beams_batch(
                 __lowercase					,      __lowercase					,      )

       SCREAMING_SNAKE_CASE__     :				List[str]									=[d[0][0] for d in decoded_decoder_out]

       self.assertListEqual(__lowercase					,      __lowercase      )
       self.assertListEqual(['''<s> </s> <s> </s> </s>''', '''</s> </s> <s> </s> </s>''']					,      __lowercase      )
       SCREAMING_SNAKE_CASE__     :				Optional[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					,      __lowercase      )




 def      __magic_name__	(      self	:							Union[str, Any]      )	->						str:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''      )

       SCREAMING_SNAKE_CASE__     :				List[str]									=processor.decoder.model_container[processor.decoder._model_key]
       SCREAMING_SNAKE_CASE__     :				Tuple									=Path(language_model._kenlm_model.path.decode('''utf-8'''      )      ).parent.parent.absolute()

       SCREAMING_SNAKE_CASE__     :				Any									=os.listdir(__lowercase      )
       SCREAMING_SNAKE_CASE__     :				List[str]									=['''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(__lowercase					,      __lowercase      )




 def      __magic_name__	(      self	:							Tuple      )	->						int:
       SCREAMING_SNAKE_CASE__     :				Any									=snapshot_download('''hf-internal-testing/processor_with_lm'''      )

       SCREAMING_SNAKE_CASE__     :				Tuple									=WavaVecaProcessorWithLM.from_pretrained(__lowercase      )

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

       SCREAMING_SNAKE_CASE__     :				int									=os.listdir(__lowercase      )
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=os.listdir(__lowercase      )

       local_decoder_files.sort()
       expected_decoder_files.sort()

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




 def      __magic_name__	(      self	:							List[str]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Dict									=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''      )
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm'''      )

       SCREAMING_SNAKE_CASE__     :				Any									=floats_list((3, 10_00)      )

       SCREAMING_SNAKE_CASE__     :				Optional[int]									=processor_wavaveca(__lowercase					,      return_tensors='''np'''      )
       SCREAMING_SNAKE_CASE__     :				List[str]									=processor_auto(__lowercase					,      return_tensors='''np'''      )

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

       SCREAMING_SNAKE_CASE__     :				Tuple									=self._get_dummy_logits()

       SCREAMING_SNAKE_CASE__     :				Any									=processor_wavaveca.batch_decode(__lowercase      )
       SCREAMING_SNAKE_CASE__     :				int									=processor_auto.batch_decode(__lowercase      )

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




 def      __magic_name__	(      self	:							Any      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				List[str]									=self.get_feature_extractor()
       SCREAMING_SNAKE_CASE__     :				str									=self.get_tokenizer()
       SCREAMING_SNAKE_CASE__     :				Any									=self.get_decoder()

       SCREAMING_SNAKE_CASE__     :				List[str]									=WavaVecaProcessorWithLM(tokenizer=__lowercase					,      feature_extractor=__lowercase					,      decoder=__lowercase      )

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




 @staticmethod
 def      __magic_name__	(      __lowercase	:							int					,      __lowercase	:							List[Any]      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				int									=[d[key] for d in offsets]
       return retrieved_list




 def      __magic_name__	(      self	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''      )
       SCREAMING_SNAKE_CASE__     :				List[Any]									=self._get_dummy_logits()[0]

       SCREAMING_SNAKE_CASE__     :				str									=processor.decode(__lowercase					,      output_word_offsets=__lowercase      )
       # 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(__lowercase					,      __lowercase      )      )

       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      __magic_name__	(      self	:							int      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''      )
       SCREAMING_SNAKE_CASE__     :				Dict									=self._get_dummy_logits()

       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=processor.batch_decode(__lowercase					,      output_word_offsets=__lowercase      )

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

       self.assertListEqual(
           [''' '''.join(self.get_from_offsets(__lowercase					,      '''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      __magic_name__	(      self	:							int      )	->						Optional[Any]:
       import torch

       SCREAMING_SNAKE_CASE__     :				int									=load_dataset('''common_voice'''					,      '''en'''					,      split='''train'''					,      streaming=__lowercase      )
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=ds.cast_column('''audio'''					,      datasets.Audio(sampling_rate=1_60_00      )      )
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=iter(__lowercase      )
       SCREAMING_SNAKE_CASE__     :				Tuple									=next(__lowercase      )

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

       with torch.no_grad():
             SCREAMING_SNAKE_CASE__     :				Any									=model(__lowercase      ).logits.cpu().numpy()

       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=processor.decode(logits[0]					,      output_word_offsets=__lowercase      )

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

       SCREAMING_SNAKE_CASE__     :				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(__lowercase					,      '''word'''      )      )					,      __lowercase      )
       self.assertEqual(''' '''.join(self.get_from_offsets(__lowercase					,      '''word'''      )      )					,      output.text      )

       # output times
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=torch.tensor(self.get_from_offsets(__lowercase					,      '''start_time'''      )      )
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=torch.tensor(self.get_from_offsets(__lowercase					,      '''end_time'''      )      )

       # fmt: off
       SCREAMING_SNAKE_CASE__     :				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]      )
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=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(__lowercase					,      __lowercase					,      atol=0.01      )      )
       self.assertTrue(torch.allclose(__lowercase					,      __lowercase					,      atol=0.01      )      )







'''simple docstring'''
import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      MegatronBertForCausalLM,
      MegatronBertForMaskedLM,
      MegatronBertForMultipleChoice,
      MegatronBertForNextSentencePrediction,
      MegatronBertForPreTraining,
      MegatronBertForQuestionAnswering,
      MegatronBertForSequenceClassification,
      MegatronBertForTokenClassification,
      MegatronBertModel,
  )






class      __SCREAMING_SNAKE_CASE      :




 def __init__(      self	:							Optional[int]					,      __lowercase	:							Optional[Any]					,      __lowercase	:							List[str]=13					,      __lowercase	:							int=7					,      __lowercase	:							Tuple=True					,      __lowercase	:							Optional[Any]=True					,      __lowercase	:							List[Any]=True					,      __lowercase	:							str=True					,      __lowercase	:							Tuple=99					,      __lowercase	:							Union[str, Any]=64					,      __lowercase	:							Tuple=32					,      __lowercase	:							Optional[Any]=5					,      __lowercase	:							Tuple=4					,      __lowercase	:							Optional[int]=37					,      __lowercase	:							int="gelu"					,      __lowercase	:							Union[str, Any]=0.1					,      __lowercase	:							List[Any]=0.1					,      __lowercase	:							List[Any]=5_12					,      __lowercase	:							int=16					,      __lowercase	:							Optional[int]=2					,      __lowercase	:							Tuple=0.02					,      __lowercase	:							List[str]=3					,      __lowercase	:							List[str]=4					,      __lowercase	:							List[str]=None					,      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=parent
       SCREAMING_SNAKE_CASE__     :				Any									=batch_size
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=seq_length
       SCREAMING_SNAKE_CASE__     :				Dict									=is_training
       SCREAMING_SNAKE_CASE__     :				List[Any]									=use_input_mask
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=use_token_type_ids
       SCREAMING_SNAKE_CASE__     :				List[Any]									=use_labels
       SCREAMING_SNAKE_CASE__     :				int									=vocab_size
       SCREAMING_SNAKE_CASE__     :				str									=hidden_size
       SCREAMING_SNAKE_CASE__     :				Any									=embedding_size
       SCREAMING_SNAKE_CASE__     :				Tuple									=num_hidden_layers
       SCREAMING_SNAKE_CASE__     :				str									=num_attention_heads
       SCREAMING_SNAKE_CASE__     :				Tuple									=intermediate_size
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=hidden_act
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=hidden_dropout_prob
       SCREAMING_SNAKE_CASE__     :				str									=attention_probs_dropout_prob
       SCREAMING_SNAKE_CASE__     :				Any									=max_position_embeddings
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=type_vocab_size
       SCREAMING_SNAKE_CASE__     :				Any									=type_sequence_label_size
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=initializer_range
       SCREAMING_SNAKE_CASE__     :				str									=num_labels
       SCREAMING_SNAKE_CASE__     :				List[str]									=num_choices
       SCREAMING_SNAKE_CASE__     :				List[str]									=scope




 def      __magic_name__	(      self	:							str      )	->						Optional[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size      )

       SCREAMING_SNAKE_CASE__     :				int									=None
       if self.use_input_mask:
             SCREAMING_SNAKE_CASE__     :				List[str]									=random_attention_mask([self.batch_size, self.seq_length]      )

       SCREAMING_SNAKE_CASE__     :				int									=None
       if self.use_token_type_ids:
             SCREAMING_SNAKE_CASE__     :				Any									=ids_tensor([self.batch_size, self.seq_length]					,      self.type_vocab_size      )

       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=None
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=None
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=None
       if self.use_labels:
             SCREAMING_SNAKE_CASE__     :				Tuple									=ids_tensor([self.batch_size]					,      self.type_sequence_label_size      )
             SCREAMING_SNAKE_CASE__     :				List[Any]									=ids_tensor([self.batch_size, self.seq_length]					,      self.num_labels      )
             SCREAMING_SNAKE_CASE__     :				Any									=ids_tensor([self.batch_size]					,      self.num_choices      )

       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.get_config()

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




 def      __magic_name__	(      self	:							List[str]      )	->						Any:
       return MegatronBertConfig(
           vocab_size=self.vocab_size					,      hidden_size=self.hidden_size					,      num_hidden_layers=self.num_hidden_layers					,      num_attention_heads=self.num_attention_heads					,      intermediate_size=self.intermediate_size					,      embedding_size=self.embedding_size					,      hidden_act=self.hidden_act					,      hidden_dropout_prob=self.hidden_dropout_prob					,      attention_probs_dropout_prob=self.attention_probs_dropout_prob					,      max_position_embeddings=self.max_position_embeddings					,      type_vocab_size=self.type_vocab_size					,      is_decoder=__lowercase					,      initializer_range=self.initializer_range					,      )




 def      __magic_name__	(      self	:							Any					,      __lowercase	:							Tuple					,      __lowercase	:							Any					,      __lowercase	:							List[Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[str]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]      )	->						int:
       SCREAMING_SNAKE_CASE__     :				List[str]									=MegatronBertModel(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[Any]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase      )
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=model(__lowercase					,      token_type_ids=__lowercase      )
       SCREAMING_SNAKE_CASE__     :				str									=model(__lowercase      )

       self.parent.assertEqual(result.last_hidden_state.shape					,      (self.batch_size, self.seq_length, self.hidden_size)      )
       self.parent.assertEqual(result.pooler_output.shape					,      (self.batch_size, self.hidden_size)      )




 def      __magic_name__	(      self	:							Dict					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[str]					,      __lowercase	:							str					,      __lowercase	:							Tuple					,      __lowercase	:							Any					,      __lowercase	:							Optional[int]					,      __lowercase	:							Any      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				str									=MegatronBertForMaskedLM(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Tuple									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)      )




 def      __magic_name__	(      self	:							List[str]					,      __lowercase	:							Tuple					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[Any]					,      __lowercase	:							Any					,      __lowercase	:							str					,      __lowercase	:							Any					,      __lowercase	:							str      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=MegatronBertForCausalLM(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Dict									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)      )




 def      __magic_name__	(      self	:							Union[str, Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							int					,      __lowercase	:							Any      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=MegatronBertForNextSentencePrediction(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, 2)      )




 def      __magic_name__	(      self	:							List[str]					,      __lowercase	:							Optional[int]					,      __lowercase	:							str					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[Any]					,      __lowercase	:							str					,      __lowercase	:							Optional[Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=MegatronBertForPreTraining(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      next_sentence_label=__lowercase					,      )
       self.parent.assertEqual(result.prediction_logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)      )
       self.parent.assertEqual(result.seq_relationship_logits.shape					,      (self.batch_size, 2)      )




 def      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[int]					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							int					,      __lowercase	:							Any					,      __lowercase	:							Any					,      __lowercase	:							List[str]      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=MegatronBertForQuestionAnswering(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Dict									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      start_positions=__lowercase					,      end_positions=__lowercase					,      )
       self.parent.assertEqual(result.start_logits.shape					,      (self.batch_size, self.seq_length)      )
       self.parent.assertEqual(result.end_logits.shape					,      (self.batch_size, self.seq_length)      )




 def      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[Any]					,      __lowercase	:							List[Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							int					,      __lowercase	:							List[str]      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.num_labels
       SCREAMING_SNAKE_CASE__     :				Dict									=MegatronBertForSequenceClassification(__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_labels)      )




 def      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							List[Any]					,      __lowercase	:							str					,      __lowercase	:							Optional[Any]					,      __lowercase	:							Any					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[Any]					,      __lowercase	:							int      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.num_labels
       SCREAMING_SNAKE_CASE__     :				int									=MegatronBertForTokenClassification(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.num_labels)      )




 def      __magic_name__	(      self	:							Dict					,      __lowercase	:							List[str]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Dict					,      __lowercase	:							Any					,      __lowercase	:							List[Any]					,      __lowercase	:							int					,      __lowercase	:							Optional[int]      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				int									=self.num_choices
       SCREAMING_SNAKE_CASE__     :				List[str]									=MegatronBertForMultipleChoice(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=input_ids.unsqueeze(1      ).expand(-1					,      self.num_choices					,      -1      ).contiguous()
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=token_type_ids.unsqueeze(1      ).expand(-1					,      self.num_choices					,      -1      ).contiguous()
       SCREAMING_SNAKE_CASE__     :				Tuple									=input_mask.unsqueeze(1      ).expand(-1					,      self.num_choices					,      -1      ).contiguous()
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_choices)      )




 def      __magic_name__	(      self	:							str      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	
       )     :				str									=config_and_inputs
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_torch
class      __SCREAMING_SNAKE_CASE      (			lowerCamelCase	,				lowerCamelCase	,				unittest.TestCase				):
 snake_case_  							=							(
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 snake_case_  							=							(
     {
         """feature-extraction""": MegatronBertModel,
         """fill-mask""": MegatronBertForMaskedLM,
         """question-answering""": MegatronBertForQuestionAnswering,
         """text-classification""": MegatronBertForSequenceClassification,
         """text-generation""": MegatronBertForCausalLM,
         """token-classification""": MegatronBertForTokenClassification,
         """zero-shot""": MegatronBertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 snake_case_  							=							True
 # test_resize_embeddings = False
 snake_case_  							=							False




 def      __magic_name__	(      self	:							List[Any]					,      __lowercase	:							List[str]					,      __lowercase	:							Tuple					,      __lowercase	:							Tuple=False      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=super()._prepare_for_class(__lowercase					,      __lowercase					,      return_labels=__lowercase      )

       if return_labels:
             if model_class in get_values(__lowercase      ):
                   SCREAMING_SNAKE_CASE__     :				List[str]									=torch.zeros(
                       (self.model_tester.batch_size, self.model_tester.seq_length)					,      dtype=torch.long					,      device=__lowercase      )
                   SCREAMING_SNAKE_CASE__     :				List[Any]									=torch.zeros(
                       self.model_tester.batch_size					,      dtype=torch.long					,      device=__lowercase      )
       return inputs_dict




 def      __magic_name__	(      self	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=MegatronBertModelTester(self      )
       SCREAMING_SNAKE_CASE__     :				Tuple									=ConfigTester(self					,      config_class=__lowercase					,      hidden_size=37      )




 def      __magic_name__	(      self	:							str      )	->						Dict:
       self.config_tester.run_common_tests()




 def      __magic_name__	(      self	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				str									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*__lowercase      )




 def      __magic_name__	(      self	:							List[str]      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				int									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__lowercase      )




 def      __magic_name__	(      self	:							Optional[Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Dict									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__lowercase      )




 def      __magic_name__	(      self	:							Optional[int]      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__lowercase      )




 def      __magic_name__	(      self	:							Dict      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*__lowercase      )




 def      __magic_name__	(      self	:							int      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*__lowercase      )




 def      __magic_name__	(      self	:							List[Any]      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__lowercase      )




 def      __magic_name__	(      self	:							Union[str, Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*__lowercase      )





def  _a(  UpperCamelCase__				:						List[str]							):




      '''simple docstring'''






      return torch.tensor(
          UpperCamelCase__,	dtype=torch.long,	device=UpperCamelCase__,	)


a_					= 1E-4







@require_torch
@require_sentencepiece
@require_tokenizers
class      __SCREAMING_SNAKE_CASE      (			unittest.TestCase				):




 @slow
 @unittest.skip('''Model is not available.'''      )
 def      __magic_name__	(      self	:							Any      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				Any									='''nvidia/megatron-bert-uncased-345m'''
       if "MYDIR" in os.environ:
             SCREAMING_SNAKE_CASE__     :				Optional[int]									=os.path.join(os.environ['''MYDIR''']					,      __lowercase      )
       SCREAMING_SNAKE_CASE__     :				Any									=MegatronBertModel.from_pretrained(__lowercase      )
       model.to(__lowercase      )
       model.half()
       SCREAMING_SNAKE_CASE__     :				Dict									=_long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      )
       with torch.no_grad():
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=model(__lowercase      )[0]
       SCREAMING_SNAKE_CASE__     :				Dict									=torch.Size((1, 9, 10_24)      )
       self.assertEqual(output.shape					,      __lowercase      )

       SCREAMING_SNAKE_CASE__     :				str									=[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
       for ii in range(3      ):
             for jj in range(3      ):
                   SCREAMING_SNAKE_CASE__     :				List[Any]									=output[0, ii, jj]
                   SCREAMING_SNAKE_CASE__     :				Tuple									=expected[3 * ii + jj]
                   SCREAMING_SNAKE_CASE__     :				List[str]									='''ii={} jj={} a={} b={}'''.format(__lowercase					,      __lowercase					,      __lowercase					,      __lowercase      )
                   self.assertTrue(math.isclose(__lowercase					,      __lowercase					,      rel_tol=__lowercase					,      abs_tol=__lowercase      )					,      msg=__lowercase      )