





"""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 		lowercase					:





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





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       ids_tensor([self.batch_size, self.seq_length]      ,						self.vocab_size )

            lowerCAmelCase  						=       None
            if self.use_input_mask:
                  lowerCAmelCase  						=       random_attention_mask([self.batch_size, self.seq_length] )

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

            lowerCAmelCase  						=       None
            lowerCAmelCase  						=       None
            lowerCAmelCase  						=       None
            if self.use_labels:
                  lowerCAmelCase  						=       ids_tensor([self.batch_size]      ,						self.type_sequence_label_size )
                  lowerCAmelCase  						=       ids_tensor([self.batch_size, self.seq_length]      ,						self.num_labels )
                  lowerCAmelCase  						=       ids_tensor([self.batch_size]      ,						self.num_choices )

            lowerCAmelCase  						=       self.get_config()

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





      def     _snake_case			(       self )		->     List[str]:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Optional[int]:
            lowerCAmelCase  						=       MegatronBertModel(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase )
            lowerCAmelCase  						=       model(lowercase      ,						token_type_ids=lowercase )
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Dict:
            lowerCAmelCase  						=       MegatronBertForMaskedLM(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Any:
            lowerCAmelCase  						=       MegatronBertForCausalLM(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     List[Any]:
            lowerCAmelCase  						=       MegatronBertForNextSentencePrediction(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(
                lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase      ,						)
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, 2) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Tuple:
            lowerCAmelCase  						=       MegatronBertForPreTraining(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     List[Any]:
            lowerCAmelCase  						=       MegatronBertForQuestionAnswering(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     List[Any]:
            lowerCAmelCase  						=       self.num_labels
            lowerCAmelCase  						=       MegatronBertForSequenceClassification(lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase )
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_labels) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.num_labels
            lowerCAmelCase  						=       MegatronBertForTokenClassification(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Optional[Any]:
            lowerCAmelCase  						=       self.num_choices
            lowerCAmelCase  						=       MegatronBertForMultipleChoice(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       input_ids.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       token_type_ids.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       input_mask.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       model(
                lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase      ,						)
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_choices) )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            (
                (
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		
            )  						=       config_and_inputs
            lowerCAmelCase  						=       {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
            return config, inputs_dict







@require_torch
class 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase     ,  unittest.TestCase					):
      _SCREAMING_SNAKE_CASE					=      (
          (
              MegatronBertModel,
              MegatronBertForMaskedLM,
              MegatronBertForCausalLM,
              MegatronBertForMultipleChoice,
              MegatronBertForNextSentencePrediction,
              MegatronBertForPreTraining,
              MegatronBertForQuestionAnswering,
              MegatronBertForSequenceClassification,
              MegatronBertForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _SCREAMING_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 {}
      )
      _SCREAMING_SNAKE_CASE					=      True
      # test_resize_embeddings = False
      _SCREAMING_SNAKE_CASE					=      False





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase=False )		->     int:
            lowerCAmelCase  						=       super()._prepare_for_class(lowercase      ,						lowercase      ,						return_labels=lowercase )

            if return_labels:
                  if model_class in get_values(lowercase ):
                        lowerCAmelCase  						=       torch.zeros(
                            (self.model_tester.batch_size, self.model_tester.seq_length)      ,						dtype=torch.long      ,						device=lowercase )
                        lowerCAmelCase  						=       torch.zeros(
                            self.model_tester.batch_size      ,						dtype=torch.long      ,						device=lowercase )
            return inputs_dict





      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       MegatronBertModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase      ,						hidden_size=37 )





      def     _snake_case			(       self )		->     Tuple:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     str:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_model(*lowercase )





      def     _snake_case			(       self )		->     List[str]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase )





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase )





      def     _snake_case			(       self )		->     Dict:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase )





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase )







def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						List[Any]			):







      '''simple docstring'''




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


SCREAMING_SNAKE_CASE__					=   1e-4






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





      @slow
      @unittest.skip("""Model is not available.""" )
      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       """nvidia/megatron-bert-uncased-345m"""
            if "MYDIR" in os.environ:
                  lowerCAmelCase  						=       os.path.join(os.environ["""MYDIR"""]      ,						lowercase )
            lowerCAmelCase  						=       MegatronBertModel.from_pretrained(lowercase )
            model.to(lowercase )
            model.half()
            lowerCAmelCase  						=       _long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]] )
            with torch.no_grad():
                  lowerCAmelCase  						=       model(lowercase )[0]
            lowerCAmelCase  						=       torch.Size((1, 9, 1_024) )
            self.assertEqual(output.shape      ,						lowercase )

            lowerCAmelCase  						=       [-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
            for ii in range(3 ):
                  for jj in range(3 ):
                        lowerCAmelCase  						=       output[0, ii, jj]
                        lowerCAmelCase  						=       expected[3 * ii + jj]
                        lowerCAmelCase  						=       """ii={} jj={} a={} b={}""".format(lowercase      ,						lowercase      ,						lowercase      ,						lowercase )
                        self.assertTrue(math.isclose(lowercase      ,						lowercase      ,						rel_tol=lowercase      ,						abs_tol=lowercase )      ,						msg=lowercase )











"""simple docstring"""
import argparse
import re
from typing import Dict

import torch
from datasets import Audio, Dataset, load_dataset, load_metric

from transformers import AutoFeatureExtractor, pipeline

def       _snake_case						(						UpperCamelCase						:		Dataset			,		UpperCamelCase						:		Dict[str, str]						):
 UpperCAmelCase	:					Any	=	args.log_outputs
 UpperCAmelCase	:					Any	=	"""_""".join(args.dataset.split("""/"""						) + [args.config, args.split]						)

 # load metric
 UpperCAmelCase	:					List[Any]	=	load_metric("""wer"""						)
 UpperCAmelCase	:					Any	=	load_metric("""cer"""						)

 # compute metrics
 UpperCAmelCase	:					int	=	wer.compute(references=result["""target"""]			,		predictions=result["""prediction"""]						)
 UpperCAmelCase	:					str	=	cer.compute(references=result["""target"""]			,		predictions=result["""prediction"""]						)

 # print & log results
 UpperCAmelCase	:					Tuple	=	F"WER: {wer_result}\nCER: {cer_result}"
 print(UpperCamelCase						)

 with open(F"{dataset_id}_eval_results.txt"			,		"""w"""						) as f:
  f.write(UpperCamelCase						)

 # log all results in text file. Possibly interesting for analysis
 if log_outputs is not None:
  UpperCAmelCase	:					str	=	F"log_{dataset_id}_predictions.txt"
  UpperCAmelCase	:					Tuple	=	F"log_{dataset_id}_targets.txt"

  with open(UpperCamelCase			,		"""w"""						) as p, open(UpperCamelCase			,		"""w"""						) as t:
   # mapping function to write output
   def write_to_file(UpperCamelCase						:		List[Any]			,		UpperCamelCase						:		List[Any]						):
    p.write(F"{i}" + """\n"""						)
    p.write(batch["""prediction"""] + """\n"""						)
    t.write(F"{i}" + """\n"""						)
    t.write(batch["""target"""] + """\n"""						)

   result.map(UpperCamelCase			,		with_indices=UpperCamelCase						)

def       _snake_case						(						UpperCamelCase						:		str						):
 UpperCAmelCase	:					List[str]	=	"""[,?.!\-\;\:\"“%‘”�—’…–]"""  # noqa: W605 IMPORTANT: this should correspond to the chars that were ignored during training

 UpperCAmelCase	:					Dict	=	re.sub(UpperCamelCase			,		""""""			,		text.lower()						)

 # In addition, we can normalize the target text, e.g. removing new lines characters etc...
 # note that order is important here!
 UpperCAmelCase	:					List[str]	=	["""\n\n""", """\n""", """   """, """  """]

 for t in token_sequences_to_ignore:
  UpperCAmelCase	:					Optional[Any]	=	""" """.join(text.split(UpperCamelCase						)						)

 return text

def       _snake_case						(						UpperCamelCase						:		Tuple						):
 # load dataset
 UpperCAmelCase	:					Union[str, Any]	=	load_dataset(args.dataset			,		args.config			,		split=args.split			,		use_auth_token=UpperCamelCase						)

 # for testing: only process the first two examples as a test
 # dataset = dataset.select(range(10))

 # load processor
 UpperCAmelCase	:					Optional[int]	=	AutoFeatureExtractor.from_pretrained(args.model_id						)
 UpperCAmelCase	:					Any	=	feature_extractor.sampling_rate

 # resample audio
 UpperCAmelCase	:					List[str]	=	dataset.cast_column("""audio"""			,		Audio(sampling_rate=UpperCamelCase						)						)

 # load eval pipeline
 if args.device is None:
  UpperCAmelCase	:					Optional[int]	=	0 if torch.cuda.is_available() else -1
 UpperCAmelCase	:					Tuple	=	pipeline("""automatic-speech-recognition"""			,		model=args.model_id			,		device=args.device						)

 # map function to decode audio
 def map_to_pred(UpperCamelCase						:		Any						):
  UpperCAmelCase	:					Any	=	asr(
      batch["""audio"""]["""array"""]			,		chunk_length_s=args.chunk_length_s			,		stride_length_s=args.stride_length_s						)

  UpperCAmelCase	:					Tuple	=	prediction["""text"""]
  UpperCAmelCase	:					List[str]	=	normalize_text(batch["""sentence"""]						)
  return batch

 # run inference on all examples
 UpperCAmelCase	:					int	=	dataset.map(UpperCamelCase			,		remove_columns=dataset.column_names						)

 # compute and log_results
 # do not change function below
 log_results(UpperCamelCase			,		UpperCamelCase						)


if __name__ == "__main__":
      A:   List[Any]								=		argparse.ArgumentParser()

      parser.add_argument(
          "--model_id", type=str, required=True, help="Model identifier. Should be loadable with 🤗 Transformers"
      )
      parser.add_argument(
          "--dataset",
          type=str,
          required=True,
          help="Dataset name to evaluate the `model_id`. Should be loadable with 🤗 Datasets",
      )
      parser.add_argument(
          "--config", type=str, required=True, help="Config of the dataset. *E.g.* `'en'`  for Common Voice"
      )
      parser.add_argument("--split", type=str, required=True, help="Split of the dataset. *E.g.* `'test'`")
      parser.add_argument(
          "--chunk_length_s", type=float, default=None, help="Chunk length in seconds. Defaults to 5 seconds."
      )
      parser.add_argument(
          "--stride_length_s", type=float, default=None, help="Stride of the audio chunks. Defaults to 1 second."
      )
      parser.add_argument(
          "--log_outputs", action="store_true", help="If defined, write outputs to log file for analysis."
      )
      parser.add_argument(
          "--device",
          type=int,
          default=None,
          help="The device to run the pipeline on. -1 for CPU (default), 0 for the first GPU and so on.",
      )
      A:   Union[str, Any]								=		parser.parse_args()

      main(args)
