

"""simple docstring"""

import argparse
import os
from pathlib import Path
from typing import Dict

import tensorflow as tf
import torch
from tqdm import tqdm

from transformers import PegasusConfig, PegasusForConditionalGeneration, PegasusTokenizer
from transformers.models.pegasus.configuration_pegasus import DEFAULTS, task_specific_params


_UpperCamelCase		     =    [
    # replace left string with right string to get the relevant state_dict key (identical state dict to bart)
    ["""memory_attention""", """encoder_attn"""],
    ["""attention""", """attn"""],
    ["""/""", """."""],
    [""".LayerNorm.gamma""", """_layer_norm.weight"""],
    [""".LayerNorm.beta""", """_layer_norm.bias"""],
    ["""r.layer_""", """r.layers."""],
    ["""output_proj""", """out_proj"""],
    ["""ffn.dense_1.""", """fc2."""],
    ["""ffn.dense.""", """fc1."""],
    ["""ffn_layer_norm""", """final_layer_norm"""],
    ["""kernel""", """weight"""],
    ["""encoder_layer_norm.""", """encoder.layer_norm."""],
    ["""decoder_layer_norm.""", """decoder.layer_norm."""],
    ["""embeddings.weights""", """shared.weight"""],
]



def       SCREAMING_SNAKE_CASE							( lowercase__   )			->						List[Any]:
 for pegasus_name, hf_name in PATTERNS:
  lowerCAmelCase__   :				Any 			=				k.replace(lowercase__					,					lowercase__   )
 return k



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__   )			->						PegasusForConditionalGeneration:
 lowerCAmelCase__   :				List[Any] 			=				DEFAULTS.copy()
 cfg_kwargs.update(lowercase__   )
 lowerCAmelCase__   :				Dict 			=				PegasusConfig(**lowercase__   )
 lowerCAmelCase__   :				int 			=				PegasusForConditionalGeneration(lowercase__   )
 lowerCAmelCase__   :				Union[str, Any] 			=				torch_model.model.state_dict()
 lowerCAmelCase__   :				Tuple 			=				{}
 for k, v in tf_weights.items():
  lowerCAmelCase__   :				List[str] 			=				rename_state_dict_key(lowercase__   )
  if new_k not in sd:
   raise ValueError(F"""could not find new key {new_k} in state dict. (converted from {k})"""   )

  if "dense" in k or "proj" in new_k:
   lowerCAmelCase__   :				Dict 			=				v.T
  lowerCAmelCase__   :				Optional[int] 			=				torch.tensor(lowercase__					,					dtype=sd[new_k].dtype   )
  assert v.shape == sd[new_k].shape, F"""{new_k}, {k}, {v.shape}, {sd[new_k].shape}"""
 # make sure embedding.padding_idx is respected
 lowerCAmelCase__   :				List[str] 			=				torch.zeros_like(mapping["shared.weight"][cfg.pad_token_id + 1]   )
 lowerCAmelCase__   :				Any 			=				mapping["shared.weight"]
 lowerCAmelCase__   :				List[str] 			=				mapping["shared.weight"]
 lowerCAmelCase__   :				Any 			=				{k: torch.zeros_like(lowercase__   ) for k, v in sd.items() if k.endswith("bias"   ) and k not in mapping}
 mapping.update(**lowercase__   )
 lowerCAmelCase__			,						lowerCAmelCase__   :				Union[str, Any] 			=				torch_model.model.load_state_dict(lowercase__					,					strict=lowercase__   )
 lowerCAmelCase__   :				Union[str, Any] 			=				[
     k for k in missing if k not in ["encoder.embed_positions.weight", "decoder.embed_positions.weight"]
 ]
 assert unexpected_missing == [], F"""no matches found for the following torch keys {unexpected_missing}"""
 assert extra == [], F"""no matches found for the following tf keys {extra}"""
 return torch_model



def       SCREAMING_SNAKE_CASE							( lowercase__="./ckpt/aeslc/model.ckpt-32000"   )			->						Dict:
 lowerCAmelCase__   :				Optional[Any] 			=				tf.train.list_variables(lowercase__   )
 lowerCAmelCase__   :				List[Any] 			=				{}
 lowerCAmelCase__   :				List[str] 			=				["Adafactor", "global_step"]
 for name, shape in tqdm(lowercase__					,					desc="converting tf checkpoint to dict"   ):
  lowerCAmelCase__   :				Dict 			=				any(pat in name for pat in ignore_name   )
  if skip_key:
   continue
  lowerCAmelCase__   :				Tuple 			=				tf.train.load_variable(lowercase__					,					lowercase__   )
  lowerCAmelCase__   :				Any 			=				array
 return tf_weights



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__   )			->						Dict:
 # save tokenizer first
 lowerCAmelCase__   :				List[str] 			=				Path(lowercase__   ).parent.name
 lowerCAmelCase__   :				List[str] 			=				task_specific_params[F"""summarization_{dataset}"""]["max_position_embeddings"]
 lowerCAmelCase__   :				Tuple 			=				PegasusTokenizer.from_pretrained("sshleifer/pegasus"					,					model_max_length=lowercase__   )
 assert tok.model_max_length == desired_max_model_length
 tok.save_pretrained(lowercase__   )

 # convert model
 lowerCAmelCase__   :				List[Any] 			=				get_tf_weights_as_numpy(lowercase__   )
 lowerCAmelCase__   :				Optional[int] 			=				task_specific_params[F"""summarization_{dataset}"""]
 if dataset == "large":
  lowerCAmelCase__   :				Tuple 			=				task_specific_params
 lowerCAmelCase__   :				List[Any] 			=				convert_pegasus(lowercase__					,					lowercase__   )
 torch_model.save_pretrained(lowercase__   )
 lowerCAmelCase__   :				Optional[int] 			=				torch_model.state_dict()
 sd.pop("model.decoder.embed_positions.weight"   )
 sd.pop("model.encoder.embed_positions.weight"   )
 torch.save(lowercase__					,					Path(lowercase__   ) / "pytorch_model.bin"   )


if __name__ == "__main__":
    _UpperCamelCase		     =    argparse.ArgumentParser()
    # Required parameters
    parser.add_argument("""tf_ckpt_path""", type=str, help="""passed to tf.train.list_variables""")
    parser.add_argument("""save_dir""", default=None, type=str, help="""Path to the output PyTorch model.""")
    _UpperCamelCase		     =    parser.parse_args()
    if args.save_dir is None:
        _UpperCamelCase		     =    Path(args.tf_ckpt_path).parent.name
        _UpperCamelCase		     =    os.path.join("""pegasus""", dataset)
    convert_pegasus_ckpt_to_pytorch(args.tf_ckpt_path, args.save_dir)









"""simple docstring"""

import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_UpperCamelCase		     =    get_tests_dir("""fixtures""")
_UpperCamelCase		     =    get_tests_dir("""fixtures/dummy_feature_extractor_config.json""")
_UpperCamelCase		     =    get_tests_dir("""fixtures/dummy-config.json""")






class 			__a					(						unittest.TestCase       ):





       """simple docstring"""


       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				0






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h"   )
        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Tuple 			=				AutoFeatureExtractor.from_pretrained(snake_case   )
        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with tempfile.TemporaryDirectory() as tmpdirname:
         lowerCAmelCase__   :				int 			=				WavaVecaConfig()

         # remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
         lowerCAmelCase__   :				Union[str, Any] 			=				AutoFeatureExtractor.from_pretrained(snake_case   ).to_dict()

         config_dict.pop("feature_extractor_type"   )
         lowerCAmelCase__   :				Any 			=				WavaVecaFeatureExtractor(**snake_case   )

         # save in new folder
         model_config.save_pretrained(snake_case   )
         config.save_pretrained(snake_case   )

         lowerCAmelCase__   :				List[str] 			=				AutoFeatureExtractor.from_pretrained(snake_case   )

         # make sure private variable is not incorrectly saved
         lowerCAmelCase__   :				Optional[int] 			=				json.loads(config.to_json_string()   )
         self.assertTrue("_processor_class" not in dict_as_saved   )

        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				AutoFeatureExtractor.from_pretrained(snake_case   )
        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaisesRegex(
            snake_case			,					"bert-base is not a local folder and is not a valid model identifier"   ):
         lowerCAmelCase__   :				Optional[int] 			=				AutoFeatureExtractor.from_pretrained("bert-base"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaisesRegex(
            snake_case			,					R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"   ):
         lowerCAmelCase__   :				Optional[Any] 			=				AutoFeatureExtractor.from_pretrained(snake_case			,					revision="aaaaaa"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaisesRegex(
            snake_case			,					"hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."			,					):
         lowerCAmelCase__   :				Dict 			=				AutoFeatureExtractor.from_pretrained("hf-internal-testing/config-no-model"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaises(snake_case   ):
         lowerCAmelCase__   :				Tuple 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"   )
        # If remote code is disabled, we can't load this config.
        with self.assertRaises(snake_case   ):
         lowerCAmelCase__   :				str 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )

        lowerCAmelCase__   :				Optional[int] 			=				AutoFeatureExtractor.from_pretrained(
            "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )
        self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )

        # Test feature extractor can be reloaded.
        with tempfile.TemporaryDirectory() as tmp_dir:
         feature_extractor.save_pretrained(snake_case   )
         lowerCAmelCase__   :				Optional[Any] 			=				AutoFeatureExtractor.from_pretrained(snake_case			,					trust_remote_code=snake_case   )
        self.assertEqual(reloaded_feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        try:
         AutoConfig.register("custom"			,					snake_case   )
         AutoFeatureExtractor.register(snake_case			,					snake_case   )
         # Trying to register something existing in the Transformers library will raise an error
         with self.assertRaises(snake_case   ):
          AutoFeatureExtractor.register(snake_case			,					snake_case   )

         # Now that the config is registered, it can be used as any other config with the auto-API
         lowerCAmelCase__   :				List[Any] 			=				CustomFeatureExtractor.from_pretrained(snake_case   )
         with tempfile.TemporaryDirectory() as tmp_dir:
          feature_extractor.save_pretrained(snake_case   )
          lowerCAmelCase__   :				Any 			=				AutoFeatureExtractor.from_pretrained(snake_case   )
          self.assertIsInstance(snake_case			,					snake_case   )

        finally:
         if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]
         if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
          del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        class 			__a					(						__magic_name__       ):





               """simple docstring"""


               __UpperCamelCase					:      Optional[int]    =						True

        try:
         AutoConfig.register("custom"			,					snake_case   )
         AutoFeatureExtractor.register(snake_case			,					snake_case   )
         # If remote code is not set, the default is to use local
         lowerCAmelCase__   :				List[Any] 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"   )
         self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )
         self.assertTrue(feature_extractor.is_local   )

         # If remote code is disabled, we load the local one.
         lowerCAmelCase__   :				Dict 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )
         self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )
         self.assertTrue(feature_extractor.is_local   )

         # If remote is enabled, we load from the Hub
         lowerCAmelCase__   :				Tuple 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )
         self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )
         self.assertTrue(not hasattr(snake_case			,					"is_local"   )   )

        finally:
         if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]
         if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
          del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]






