





import shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
 from transformers import WhisperFeatureExtractor, WhisperProcessor


_lowerCAmelCase						=	50_358
_lowerCAmelCase						=	50_362


@require_torch
@require_torchaudio
@require_sentencepiece
class 					__magic_name__  (							unittest.TestCase							):





     """simple docstring"""



     def 						SCREAMING_SNAKE_CASE		(							self       :Union[str, Any]  ):





        '''simple docstring'''




        A_					:				List[str]          =			"openai/whisper-small.en"
        A_					:				Tuple          =			tempfile.mkdtemp()



     def 						SCREAMING_SNAKE_CASE		(							self       :Optional[Any]			,      **snake_case       :Optional[Any]  ):





        '''simple docstring'''




        return WhisperTokenizer.from_pretrained(self.checkpoint			,      **snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Dict			,      **snake_case       :Optional[int]  ):





        '''simple docstring'''




        return WhisperFeatureExtractor.from_pretrained(self.checkpoint			,      **snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :List[Any]  ):





        '''simple docstring'''




        shutil.rmtree(self.tmpdirname  )



     def 						SCREAMING_SNAKE_CASE		(							self       :List[str]  ):





        '''simple docstring'''




        A_					:				str          =			self.get_tokenizer()
        A_					:				Dict          =			self.get_feature_extractor()

        A_					:				Any          =			WhisperProcessor(tokenizer=snake_case			,      feature_extractor=snake_case  )

        processor.save_pretrained(self.tmpdirname  )
        A_					:				List[Any]          =			WhisperProcessor.from_pretrained(self.tmpdirname  )

        self.assertEqual(processor.tokenizer.get_vocab()			,      tokenizer.get_vocab()  )
        self.assertIsInstance(processor.tokenizer			,      snake_case  )

        self.assertEqual(processor.feature_extractor.to_json_string()			,      feature_extractor.to_json_string()  )
        self.assertIsInstance(processor.feature_extractor			,      snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :int  ):





        '''simple docstring'''




        A_					:				Tuple          =			WhisperProcessor(tokenizer=self.get_tokenizer()			,      feature_extractor=self.get_feature_extractor()  )
        processor.save_pretrained(self.tmpdirname  )

        A_					:				Union[str, Any]          =			self.get_tokenizer(bos_token="(BOS)"			,      eos_token="(EOS)"  )
        A_					:				List[str]          =			self.get_feature_extractor(do_normalize=snake_case			,      padding_value=1.0  )

        A_					:				Union[str, Any]          =			WhisperProcessor.from_pretrained(
            self.tmpdirname			,      bos_token="(BOS)"			,      eos_token="(EOS)"			,      do_normalize=snake_case			,      padding_value=1.0  )

        self.assertEqual(processor.tokenizer.get_vocab()			,      tokenizer_add_kwargs.get_vocab()  )
        self.assertIsInstance(processor.tokenizer			,      snake_case  )

        self.assertEqual(processor.feature_extractor.to_json_string()			,      feature_extractor_add_kwargs.to_json_string()  )
        self.assertIsInstance(processor.feature_extractor			,      snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :str  ):





        '''simple docstring'''




        A_					:				List[str]          =			self.get_feature_extractor()
        A_					:				Optional[Any]          =			self.get_tokenizer()

        A_					:				int          =			WhisperProcessor(tokenizer=snake_case			,      feature_extractor=snake_case  )

        A_					:				Any          =			floats_list((3, 1_000)  )

        A_					:				Dict          =			feature_extractor(snake_case			,      return_tensors="np"  )
        A_					:				Dict          =			processor(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 						SCREAMING_SNAKE_CASE		(							self       :Optional[Any]  ):





        '''simple docstring'''




        A_					:				List[str]          =			self.get_feature_extractor()
        A_					:				str          =			self.get_tokenizer()

        A_					:				Dict          =			WhisperProcessor(tokenizer=snake_case			,      feature_extractor=snake_case  )

        A_					:				Any          =			"This is a test string"

        A_					:				str          =			processor(text=snake_case  )

        A_					:				Union[str, Any]          =			tokenizer(snake_case  )

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



     def 						SCREAMING_SNAKE_CASE		(							self       :Union[str, Any]  ):





        '''simple docstring'''




        A_					:				List[str]          =			self.get_feature_extractor()
        A_					:				List[Any]          =			self.get_tokenizer()

        A_					:				Optional[Any]          =			WhisperProcessor(tokenizer=snake_case			,      feature_extractor=snake_case  )

        A_					:				int          =			[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        A_					:				Optional[int]          =			processor.batch_decode(snake_case  )
        A_					:				Optional[int]          =			tokenizer.batch_decode(snake_case  )

        self.assertListEqual(snake_case			,      snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :str  ):





        '''simple docstring'''




        A_					:				Optional[Any]          =			self.get_feature_extractor()
        A_					:				Union[str, Any]          =			self.get_tokenizer()

        A_					:				Dict          =			WhisperProcessor(tokenizer=snake_case			,      feature_extractor=snake_case  )

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



     def 						SCREAMING_SNAKE_CASE		(							self       :Tuple  ):





        '''simple docstring'''




        A_					:				Dict          =			self.get_feature_extractor()
        A_					:				Any          =			self.get_tokenizer()

        A_					:				Dict          =			WhisperProcessor(tokenizer=snake_case			,      feature_extractor=snake_case  )
        A_					:				Tuple          =			processor.get_decoder_prompt_ids(task="transcribe"			,      no_timestamps=snake_case  )

        self.assertIsInstance(snake_case			,      snake_case  )
        for ids in forced_decoder_ids:
           self.assertIsInstance(snake_case			,      (list, tuple)  )

        A_					:				int          =			[TRANSCRIBE, NOTIMESTAMPS]
        self.assertListEqual([ids[-1] for ids in forced_decoder_ids]			,      snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Dict  ):





        '''simple docstring'''




        A_					:				List[str]          =			WhisperProcessor(tokenizer=self.get_tokenizer()			,      feature_extractor=self.get_feature_extractor()  )
        A_					:				Any          =			processor.get_prompt_ids("Mr. Quilter"  )
        A_					:				int          =			processor.tokenizer.decode(snake_case  )

        self.assertListEqual(prompt_ids.tolist()			,      [50_360, 1_770, 13, 2_264, 346, 353]  )
        self.assertEqual(snake_case			,      "<|startofprev|> Mr. Quilter"  )



     def 						SCREAMING_SNAKE_CASE		(							self       :List[Any]  ):





        '''simple docstring'''




        A_					:				Optional[Any]          =			WhisperProcessor(tokenizer=self.get_tokenizer()			,      feature_extractor=self.get_feature_extractor()  )
        A_					:				Optional[Any]          =			processor.get_prompt_ids(""  )
        A_					:				Tuple          =			processor.tokenizer.decode(snake_case  )

        self.assertListEqual(prompt_ids.tolist()			,      [50_360, 220]  )
        self.assertEqual(snake_case			,      "<|startofprev|> "  )





     def 						SCREAMING_SNAKE_CASE		(							self       :Optional[Any]  ):





        '''simple docstring'''




        A_					:				Optional[Any]          =			WhisperProcessor(tokenizer=self.get_tokenizer()			,      feature_extractor=self.get_feature_extractor()  )

        def _test_prompt_error_raised_helper(snake_case       :List[Any]			,      snake_case       :Optional[int]  ):
           with pytest.raises(snake_case  ) as excinfo:
              processor.get_prompt_ids(snake_case  )
           A_					:				Any          =			f"Encountered text in the prompt corresponding to disallowed special token: {special_token}."
           self.assertEqual(snake_case			,      str(excinfo.value  )  )

        _test_prompt_error_raised_helper("<|startofprev|> test"			,      "<|startofprev|>"  )
        _test_prompt_error_raised_helper("test <|notimestamps|>"			,      "<|notimestamps|>"  )
        _test_prompt_error_raised_helper("test <|zh|> test <|transcribe|>"			,      "<|zh|>"  )







import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()

def      __snake_case						(					_lowerCAmelCase : Optional[int]				,				_lowerCAmelCase : List[str]				,				_lowerCAmelCase : List[Any]	) ->						Any:
   # Initialise PyTorch model
   A_					:				Tuple          =			FNetConfig.from_json_file(_lowerCAmelCase	)
   print(f"Building PyTorch model from configuration: {config}"	)
   A_					:				Tuple          =			FNetForPreTraining(_lowerCAmelCase	)

   A_					:				List[str]          =			restore_checkpoint(_lowerCAmelCase				,				_lowerCAmelCase	)
   A_					:				Any          =			checkpoint_dict["target"]

   # Embeddings
   # Position IDs
   A_					:				Optional[Any]          =			fnet_pretraining_model.state_dict()

   A_					:				Optional[int]          =			state_dict["fnet.embeddings.position_ids"]
   A_					:				Union[str, Any]          =			{"fnet.embeddings.position_ids": position_ids}
   # Embedding Layers
   A_					:				str          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]	)
   A_					:				List[Any]          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["position"]["embedding"][0]	)
   A_					:				int          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["type"]["embedding"]	)
   A_					:				List[Any]          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["kernel"]	).T
   A_					:				Dict          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["bias"]	)
   A_					:				int          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["layer_norm"]["scale"]	)
   A_					:				Any          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["layer_norm"]["bias"]	)

   # Encoder Layers
   for layer in range(config.num_hidden_layers	):
      A_					:				List[str]          =			torch.tensor(
          pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["scale"]	)
      A_					:				Union[str, Any]          =			torch.tensor(
          pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["bias"]	)

      A_					:				Union[str, Any]          =			torch.tensor(
          pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["kernel"]	).T
      A_					:				List[Any]          =			torch.tensor(
          pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["bias"]	)

      A_					:				Optional[int]          =			torch.tensor(
          pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["kernel"]	).T
      A_					:				int          =			torch.tensor(
          pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["bias"]	)

      A_					:				Any          =			torch.tensor(
          pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["scale"]	)
      A_					:				Tuple          =			torch.tensor(
          pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["bias"]	)

   # Pooler Layers
   A_					:				Optional[Any]          =			torch.tensor(pretrained_model_params["encoder"]["pooler"]["kernel"]	).T
   A_					:				List[Any]          =			torch.tensor(pretrained_model_params["encoder"]["pooler"]["bias"]	)

   # Masked LM Layers
   A_					:				Optional[Any]          =			torch.tensor(
       pretrained_model_params["predictions_dense"]["kernel"]	).T
   A_					:				int          =			torch.tensor(
       pretrained_model_params["predictions_dense"]["bias"]	)
   A_					:				str          =			torch.tensor(
       pretrained_model_params["predictions_layer_norm"]["scale"]	)
   A_					:				List[Any]          =			torch.tensor(
       pretrained_model_params["predictions_layer_norm"]["bias"]	)
   A_					:				Tuple          =			torch.tensor(
       pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]	)
   A_					:				Any          =			torch.tensor(
       pretrained_model_params["predictions_output"]["output_bias"]	)
   A_					:				List[str]          =			torch.tensor(pretrained_model_params["predictions_output"]["output_bias"]	)

   # Seq Relationship Layers
   A_					:				str          =			torch.tensor(
       pretrained_model_params["classification"]["output_kernel"]	)
   A_					:				str          =			torch.tensor(
       pretrained_model_params["classification"]["output_bias"]	)

   # Load State Dict
   fnet_pretraining_model.load_state_dict(_lowerCAmelCase	)

   # Save PreTrained
   print(f"Saving pretrained model to {save_path}"	)
   fnet_pretraining_model.save_pretrained(_lowerCAmelCase	)


if __name__ == "__main__":
 _lowerCAmelCase							:      Tuple						=	argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--flax_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
 )
 parser.add_argument(
     '''--fnet_config_file''',
     default=None,
     type=str,
     required=True,
     help=(
         '''The config json file corresponding to the pre-trained FNet model. \n'''
         '''This specifies the model architecture.'''
     ),
 )
 parser.add_argument('''--save_path''', default=None, type=str, required=True, help='''Path to the output model.''')
 _lowerCAmelCase							:      Optional[int]						=	parser.parse_args()
 convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)
