





from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


__lowerCAmelCase 	=      logging.get_logger(__name__)

__lowerCAmelCase 	=      {
    "gpt2": "https://huggingface.co/gpt2/resolve/main/config.json",
    "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/config.json",
    "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/config.json",
    "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/config.json",
    "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/config.json",
}




class 	_UpperCAmelCase						( __UpperCAmelCase       ):
					__A        = """gpt2"""
					__A        = ["""past_key_values"""]
					__A        = {
					    """hidden_size""": """n_embd""",
					    """max_position_embeddings""": """n_positions""",
					    """num_attention_heads""": """n_head""",
					    """num_hidden_layers""": """n_layer""",
					}




					def __init__(    self     ,   __lowerCAmelCase=50_257     ,   __lowerCAmelCase=1_024     ,   __lowerCAmelCase=768     ,   __lowerCAmelCase=12     ,   __lowerCAmelCase=12     ,   __lowerCAmelCase=None     ,   __lowerCAmelCase="gelu_new"     ,   __lowerCAmelCase=0.1     ,   __lowerCAmelCase=0.1     ,   __lowerCAmelCase=0.1     ,   __lowerCAmelCase=1E-5     ,   __lowerCAmelCase=0.0_2     ,   __lowerCAmelCase="cls_index"     ,   __lowerCAmelCase=True     ,   __lowerCAmelCase=None     ,   __lowerCAmelCase=True     ,   __lowerCAmelCase=0.1     ,   __lowerCAmelCase=True     ,   __lowerCAmelCase=True     ,   __lowerCAmelCase=50_256     ,   __lowerCAmelCase=50_256     ,   __lowerCAmelCase=False     ,   __lowerCAmelCase=False     ,   **__lowerCAmelCase     ,   )    ->	Tuple:
							UpperCamelCase__								=						vocab_size
							UpperCamelCase__								=						n_positions
							UpperCamelCase__								=						n_embd
							UpperCamelCase__								=						n_layer
							UpperCamelCase__								=						n_head
							UpperCamelCase__								=						n_inner
							UpperCamelCase__								=						activation_function
							UpperCamelCase__								=						resid_pdrop
							UpperCamelCase__								=						embd_pdrop
							UpperCamelCase__								=						attn_pdrop
							UpperCamelCase__								=						layer_norm_epsilon
							UpperCamelCase__								=						initializer_range
							UpperCamelCase__								=						summary_type
							UpperCamelCase__								=						summary_use_proj
							UpperCamelCase__								=						summary_activation
							UpperCamelCase__								=						summary_first_dropout
							UpperCamelCase__								=						summary_proj_to_labels
							UpperCamelCase__								=						scale_attn_weights
							UpperCamelCase__								=						use_cache
							UpperCamelCase__								=						scale_attn_by_inverse_layer_idx
							UpperCamelCase__								=						reorder_and_upcast_attn

							UpperCamelCase__								=						bos_token_id
							UpperCamelCase__								=						eos_token_id

							super().__init__(bos_token_id=UpperCAmelCase_     ,   eos_token_id=UpperCAmelCase_     ,   **UpperCAmelCase_							)




class 	_UpperCAmelCase						( __UpperCAmelCase       ):




					def __init__(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase = "default"     ,   __lowerCAmelCase = None     ,   __lowerCAmelCase = False     ,   )    ->	Tuple:
							super().__init__(UpperCAmelCase_     ,   task=UpperCAmelCase_     ,   patching_specs=UpperCAmelCase_     ,   use_past=UpperCAmelCase_							)
							if not getattr(self._config     ,   """pad_token_id"""     ,   UpperCAmelCase_							):
									# TODO: how to do that better?
									UpperCamelCase__								=						0




					@property
					def 						snake_case_		(    self							)    ->	Dict:
							UpperCamelCase__								=						OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}							)
							if self.use_past:
									self.fill_with_past_key_values_(UpperCAmelCase_     ,   direction="""inputs"""							)
									UpperCamelCase__								=						{0: """batch""", 1: """past_sequence + sequence"""}
							else:
									UpperCamelCase__								=						{0: """batch""", 1: """sequence"""}

							return common_inputs




					@property
					def 						snake_case_		(    self							)    ->	List[str]:
							return self._config.n_layer




					@property
					def 						snake_case_		(    self							)    ->	int:
							return self._config.n_head




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase = -1     ,   __lowerCAmelCase = -1     ,   __lowerCAmelCase = False     ,   __lowerCAmelCase = None     ,   )    ->	Dict:
							UpperCamelCase__								=						super(UpperCAmelCase_     ,   self							).generate_dummy_inputs(
							    UpperCAmelCase_     ,   batch_size=UpperCAmelCase_     ,   seq_length=UpperCAmelCase_     ,   is_pair=UpperCAmelCase_     ,   framework=UpperCAmelCase_							)

							# We need to order the input in the way they appears in the forward()
							UpperCamelCase__								=						OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}							)

							# Need to add the past_keys
							if self.use_past:
									if not is_torch_available():
											raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""							)
									else:
											import torch

											UpperCamelCase__	,      UpperCamelCase__								=						common_inputs["""input_ids"""].shape
											# Not using the same length for past_key_values
											UpperCamelCase__								=						seqlen + 2
											UpperCamelCase__								=						(
											    batch,
											    self.num_attention_heads,
											    past_key_values_length,
											    self._config.hidden_size // self.num_attention_heads,
											)
											UpperCamelCase__								=						[
											    (torch.zeros(UpperCAmelCase_							), torch.zeros(UpperCAmelCase_							)) for _ in range(self.num_layers							)
											]

							UpperCamelCase__								=						common_inputs["""attention_mask"""]
							if self.use_past:
									UpperCamelCase__								=						ordered_inputs["""attention_mask"""].dtype
									UpperCamelCase__								=						torch.cat(
									    [ordered_inputs["""attention_mask"""], torch.ones(UpperCAmelCase_     ,   UpperCAmelCase_     ,   dtype=UpperCAmelCase_							)]     ,   dim=1							)

							return ordered_inputs




					@property
					def 						snake_case_		(    self							)    ->	Optional[int]:
							return 13











import argparse

from transformers import CanineConfig, CanineModel, CanineTokenizer, load_tf_weights_in_canine
from transformers.utils import logging


logging.set_verbosity_info()



def    UpperCAmelCase							(  _snake_case    ,			_snake_case				):
							# Initialize PyTorch model
							lowerCAmelCase =		CanineConfig()
							lowerCAmelCase =		CanineModel(_snake_case				)
							model.eval()

							print(F"""Building PyTorch model from configuration: {config}"""				)

							# Load weights from tf checkpoint
							load_tf_weights_in_canine(_snake_case    ,			_snake_case    ,			_snake_case				)

							# Save pytorch-model (weights and configuration)
							print(F"""Save PyTorch model to {pytorch_dump_path}"""				)
							model.save_pretrained(_snake_case				)

							# Save tokenizer files
							lowerCAmelCase =		CanineTokenizer()
							print(F"""Save tokenizer files to {pytorch_dump_path}"""				)
							tokenizer.save_pretrained(_snake_case				)


if __name__ == "__main__":
		UpperCAmelCase_    =argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--tf_checkpoint_path""",
		    default=None,
		    type=str,
		    required=True,
		    help="""Path to the TensorFlow checkpoint. Should end with model.ckpt""",
		)
		parser.add_argument(
		    """--pytorch_dump_path""",
		    default=None,
		    type=str,
		    required=True,
		    help="""Path to a folder where the PyTorch model will be placed.""",
		)
		UpperCAmelCase_    =parser.parse_args()
		convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.pytorch_dump_path)



