
"""simple docstring"""


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

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402







class 				lowerCAmelCase_				(      unittest.TestCase    ):



						'''simple docstring'''


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any      )     ->   Any:
							A						   =   0


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any      )     ->   Optional[int]:
							A						   =   AutoImageProcessor.from_pretrained('openai/clip-vit-base-patch32'      )
							self.assertIsInstance(A_  ,A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any      )     ->   Any:
							with tempfile.TemporaryDirectory() as tmpdirname:
								A						   =   Path(A_      ) / 'preprocessor_config.json'
								A						   =   Path(A_      ) / 'config.json'
								json.dump(
								    {'image_processor_type': 'CLIPImageProcessor', 'processor_class': 'CLIPProcessor'}  ,open(A_  ,'w'      )  ,)
								json.dump({'model_type': 'clip'}  ,open(A_  ,'w'      )      )

								A						   =   AutoImageProcessor.from_pretrained(A_      )
								self.assertIsInstance(A_  ,A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   Union[str, Any]:
							# Ensure we can load the image processor from the feature extractor config
							with tempfile.TemporaryDirectory() as tmpdirname:
								A						   =   Path(A_      ) / 'preprocessor_config.json'
								A						   =   Path(A_      ) / 'config.json'
								json.dump(
								    {'feature_extractor_type': 'CLIPFeatureExtractor', 'processor_class': 'CLIPProcessor'}  ,open(A_  ,'w'      )  ,)
								json.dump({'model_type': 'clip'}  ,open(A_  ,'w'      )      )

								A						   =   AutoImageProcessor.from_pretrained(A_      )
								self.assertIsInstance(A_  ,A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]      )     ->   Union[str, Any]:
							with tempfile.TemporaryDirectory() as tmpdirname:
								A						   =   CLIPConfig()

								# Create a dummy config file with image_proceesor_type
								A						   =   Path(A_      ) / 'preprocessor_config.json'
								A						   =   Path(A_      ) / 'config.json'
								json.dump(
								    {'image_processor_type': 'CLIPImageProcessor', 'processor_class': 'CLIPProcessor'}  ,open(A_  ,'w'      )  ,)
								json.dump({'model_type': 'clip'}  ,open(A_  ,'w'      )      )

								# remove image_processor_type to make sure config.json alone is enough to load image processor locally
								A						   =   AutoImageProcessor.from_pretrained(A_      ).to_dict()

								config_dict.pop('image_processor_type'      )
								A						   =   CLIPImageProcessor(**A_      )

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

								A						   =   AutoImageProcessor.from_pretrained(A_      )

								# make sure private variable is not incorrectly saved
								A						   =   json.loads(config.to_json_string()      )
								self.assertTrue('_processor_class' not in dict_as_saved      )

							self.assertIsInstance(A_  ,A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Union[str, Any]:
							with tempfile.TemporaryDirectory() as tmpdirname:
								A						   =   Path(A_      ) / 'preprocessor_config.json'
								json.dump(
								    {'image_processor_type': 'CLIPImageProcessor', 'processor_class': 'CLIPProcessor'}  ,open(A_  ,'w'      )  ,)

								A						   =   AutoImageProcessor.from_pretrained(A_      )
								self.assertIsInstance(A_  ,A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   Optional[Any]:
							with self.assertRaisesRegex(
							    A_  ,'clip-base is not a local folder and is not a valid model identifier'      ):
								A						   =   AutoImageProcessor.from_pretrained('clip-base'      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   Dict:
							with self.assertRaisesRegex(
							    A_  ,R'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'      ):
								A						   =   AutoImageProcessor.from_pretrained(A_  ,revision='aaaaaa'      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   int:
							with self.assertRaisesRegex(
							    A_  ,'hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'  ,):
								A						   =   AutoImageProcessor.from_pretrained('hf-internal-testing/config-no-model'      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   int:
							# If remote code is not set, we will time out when asking whether to load the model.
							with self.assertRaises(A_      ):
								A						   =   AutoImageProcessor.from_pretrained('hf-internal-testing/test_dynamic_image_processor'      )
							# If remote code is disabled, we can't load this config.
							with self.assertRaises(A_      ):
								A						   =   AutoImageProcessor.from_pretrained(
								    'hf-internal-testing/test_dynamic_image_processor'  ,trust_remote_code=A_      )

							A						   =   AutoImageProcessor.from_pretrained(
							    'hf-internal-testing/test_dynamic_image_processor'  ,trust_remote_code=A_      )
							self.assertEqual(image_processor.__class__.__name__  ,'NewImageProcessor'      )

							# Test image processor can be reloaded.
							with tempfile.TemporaryDirectory() as tmp_dir:
								image_processor.save_pretrained(A_      )
								A						   =   AutoImageProcessor.from_pretrained(A_  ,trust_remote_code=A_      )
							self.assertEqual(reloaded_image_processor.__class__.__name__  ,'NewImageProcessor'      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   List[str]:
							try:
								AutoConfig.register('custom'  ,A_      )
								AutoImageProcessor.register(A_  ,A_      )
								# Trying to register something existing in the Transformers library will raise an error
								with self.assertRaises(A_      ):
									AutoImageProcessor.register(A_  ,A_      )

								with tempfile.TemporaryDirectory() as tmpdirname:
									A						   =   Path(A_      ) / 'preprocessor_config.json'
									A						   =   Path(A_      ) / 'config.json'
									json.dump(
									    {'feature_extractor_type': 'CLIPFeatureExtractor', 'processor_class': 'CLIPProcessor'}  ,open(A_  ,'w'      )  ,)
									json.dump({'model_type': 'clip'}  ,open(A_  ,'w'      )      )

									A						   =   CustomImageProcessor.from_pretrained(A_      )

								# Now that the config is registered, it can be used as any other config with the auto-API
								with tempfile.TemporaryDirectory() as tmp_dir:
									image_processor.save_pretrained(A_      )
									A						   =   AutoImageProcessor.from_pretrained(A_      )
									self.assertIsInstance(A_  ,A_      )

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




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Optional[Any]:







							class 				lowerCAmelCase_				(      _lowercase    ):



													'''simple docstring'''


													_lowerCamelCase:    Optional[Any]      =    True

							try:
								AutoConfig.register('custom'  ,A_      )
								AutoImageProcessor.register(A_  ,A_      )
								# If remote code is not set, the default is to use local
								A						   =   AutoImageProcessor.from_pretrained('hf-internal-testing/test_dynamic_image_processor'      )
								self.assertEqual(image_processor.__class__.__name__  ,'NewImageProcessor'      )
								self.assertTrue(image_processor.is_local      )

								# If remote code is disabled, we load the local one.
								A						   =   AutoImageProcessor.from_pretrained(
								    'hf-internal-testing/test_dynamic_image_processor'  ,trust_remote_code=A_      )
								self.assertEqual(image_processor.__class__.__name__  ,'NewImageProcessor'      )
								self.assertTrue(image_processor.is_local      )

								# If remote is enabled, we load from the Hub
								A						   =   AutoImageProcessor.from_pretrained(
								    'hf-internal-testing/test_dynamic_image_processor'  ,trust_remote_code=A_      )
								self.assertEqual(image_processor.__class__.__name__  ,'NewImageProcessor'      )
								self.assertTrue(not hasattr(A_  ,'is_local'      )      )

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

"""simple docstring"""


import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_lowercase							   =				data_utils.TransfoXLTokenizer
_lowercase							   =				data_utils.TransfoXLCorpus
_lowercase							   =				data_utils
_lowercase							   =				data_utils






def       _snake_case				(		snake_case__   :     Union[str, Any]	,      snake_case__   :     Dict	,      snake_case__   :     Tuple	,      snake_case__   :     int     ):
	if transfo_xl_dataset_file:
		# Convert a pre-processed corpus (see original TensorFlow repo)
		with open(snake_case__	,      'rb'     ) as fp:
			A						   =   pickle.load(snake_case__	,      encoding='latin1'     )
		# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
		A						   =   pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
		print(F'Save vocabulary to {pytorch_vocab_dump_path}'     )
		A						   =   corpus.vocab.__dict__
		torch.save(snake_case__	,      snake_case__     )

		A						   =   corpus.__dict__
		corpus_dict_no_vocab.pop('vocab'	,      snake_case__     )
		A						   =   pytorch_dump_folder_path + '/' + CORPUS_NAME
		print(F'Save dataset to {pytorch_dataset_dump_path}'     )
		torch.save(snake_case__	,      snake_case__     )

	if tf_checkpoint_path:
		# Convert a pre-trained TensorFlow model
		A						   =   os.path.abspath(snake_case__     )
		A						   =   os.path.abspath(snake_case__     )

		print(F'Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.'     )
		# Initialise PyTorch model
		if transfo_xl_config_file == "":
			A						   =   TransfoXLConfig()
		else:
			A						   =   TransfoXLConfig.from_json_file(snake_case__     )
		print(F'Building PyTorch model from configuration: {config}'     )
		A						   =   TransfoXLLMHeadModel(snake_case__     )

		A						   =   load_tf_weights_in_transfo_xl(snake_case__	,      snake_case__	,      snake_case__     )
		# Save pytorch-model
		A						   =   os.path.join(snake_case__	,      snake_case__     )
		A						   =   os.path.join(snake_case__	,      snake_case__     )
		print(F'Save PyTorch model to {os.path.abspath(snake_case__     )}'     )
		torch.save(model.state_dict()	,      snake_case__     )
		print(F'Save configuration file to {os.path.abspath(snake_case__     )}'     )
		with open(snake_case__	,      'w'	,      encoding='utf-8'     ) as f:
			f.write(config.to_json_string()     )


if __name__ == "__main__":
		_lowercase							   =				argparse.ArgumentParser()
		parser.add_argument(
		    '''--pytorch_dump_folder_path''',
		    default=None,
		    type=str,
		    required=True,
		    help='''Path to the folder to store the PyTorch model or dataset/vocab.''',
		)
		parser.add_argument(
		    '''--tf_checkpoint_path''',
		    default='''''',
		    type=str,
		    help='''An optional path to a TensorFlow checkpoint path to be converted.''',
		)
		parser.add_argument(
		    '''--transfo_xl_config_file''',
		    default='''''',
		    type=str,
		    help=(
		        '''An optional config json file corresponding to the pre-trained BERT model. \n'''
		        '''This specifies the model architecture.'''
		    ),
		)
		parser.add_argument(
		    '''--transfo_xl_dataset_file''',
		    default='''''',
		    type=str,
		    help='''An optional dataset file to be converted in a vocabulary.''',
		)
		_lowercase							   =				parser.parse_args()
		convert_transfo_xl_checkpoint_to_pytorch(
		    args.tf_checkpoint_path,
		    args.transfo_xl_config_file,
		    args.pytorch_dump_folder_path,
		    args.transfo_xl_dataset_file,
		)