





import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import torch
from datasets import load_dataset
from torchvision.transforms import Compose, Lambda, Normalize, RandomHorizontalFlip, RandomResizedCrop, ToTensor
from torchvision.transforms.functional import InterpolationMode

import transformers
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    ViTImageProcessor,
    ViTMAEConfig,
    ViTMAEForPreTraining,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
a__:      Optional[int]									=		logging.getLogger(__name__)

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version('4.31.0')

require_version('datasets>=1.8.0', 'To fix: pip install -r examples/pytorch/image-pretraining/requirements.txt')




@dataclass
class     SCREAMING_SNAKE_CASE__ :
 __SCREAMING_SNAKE_CASE							       =	field(
     default='''cifar10'''				,					metadata={'''help''': '''Name of a dataset from the datasets package'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={'''help''': '''The configuration name of the dataset to use (via the datasets library).'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={'''help''': '''The column name of the images in the files.'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(default=UpperCamelCase__				,					metadata={'''help''': '''A folder containing the training data.'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(default=UpperCamelCase__				,					metadata={'''help''': '''A folder containing the validation data.'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default=0.1_5				,					metadata={'''help''': '''Percent to split off of train for validation.'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={
         '''help''': (
             '''For debugging purposes or quicker training, truncate the number of training examples to this '''
             '''value if set.'''
         )
     }				,					)
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={
         '''help''': (
             '''For debugging purposes or quicker training, truncate the number of evaluation examples to this '''
             '''value if set.'''
         )
     }				,					)



 def 				UpperCamelCase			(							self ):
        A__											=						{}
        if self.train_dir is not None:
               A__											=						self.train_dir
        if self.validation_dir is not None:
               A__											=						self.validation_dir
        A__											=						data_files if data_files else None




@dataclass
class     SCREAMING_SNAKE_CASE__ :
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={
         '''help''': (
             '''The model checkpoint for weights initialization.Don\'t set if you want to train a model from scratch.'''
         )
     }				,					)
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={'''help''': '''Pretrained config name or path if not the same as model_name_or_path'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={
         '''help''': (
             '''Override some existing default config settings when a model is trained from scratch. Example: '''
             '''n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index'''
         )
     }				,					)
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={'''help''': '''Where do you want to store the pretrained models downloaded from s3'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default='''main'''				,					metadata={'''help''': '''The specific model version to use (can be a branch name, tag name or commit id).'''}				,					)
 __SCREAMING_SNAKE_CASE							       =	field(default=UpperCamelCase__				,					metadata={'''help''': '''Name or path of preprocessor config.'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={
         '''help''': (
             '''Will use the token generated when running `huggingface-cli login` (necessary to use this script '''
             '''with private models).'''
         )
     }				,					)
 __SCREAMING_SNAKE_CASE							       =	field(
     default=0.7_5				,					metadata={'''help''': '''The ratio of the number of masked tokens in the input sequence.'''}      )
 __SCREAMING_SNAKE_CASE							       =	field(
     default=UpperCamelCase__				,					metadata={'''help''': '''Whether or not to train with normalized pixel values as target.'''}      )




@dataclass
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):
 __SCREAMING_SNAKE_CASE							       =	field(
     default=1E-3				,					metadata={'''help''': '''Base learning rate: absolute_lr = base_lr * total_batch_size / 256.'''}      )


def 	UpperCamelCase__(       UpperCamelCase__			:       List[str]  )->Union[str, Any]:
       A__											=						torch.stack([example['''pixel_values'''] for example in examples]  )
       return {"pixel_values": pixel_values}


def 	UpperCamelCase__(       )->List[Any]:
       # See all possible arguments in src/transformers/training_args.py
       # or by passing the --help flag to this script.
       # We now keep distinct sets of args, for a cleaner separation of concerns.

       A__											=						HfArgumentParser((ModelArguments, DataTrainingArguments, CustomTrainingArguments)  )
       if len(sys.argv  ) == 2 and sys.argv[1].endswith('''.json'''  ):
              # If we pass only one argument to the script and it's the path to a json file,
              # let's parse it to get our arguments.
              A__					,	A__					,	A__											=						parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]  )  )
       else:
              A__					,	A__					,	A__											=						parser.parse_args_into_dataclasses()

       # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
       # information sent is the one passed as arguments along with your Python/PyTorch versions.
       send_example_telemetry('''run_mae''' ,     UpperCamelCase__ ,     UpperCamelCase__  )

       # Setup logging
       logging.basicConfig(
           format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''' ,     datefmt='''%m/%d/%Y %H:%M:%S''' ,     handlers=[logging.StreamHandler(sys.stdout  )] ,     )

       if training_args.should_log:
              # The default of training_args.log_level is passive, so we set log level at info here to have that default.
              transformers.utils.logging.set_verbosity_info()

       A__											=						training_args.get_process_log_level()
       logger.setLevel(UpperCamelCase__  )
       transformers.utils.logging.set_verbosity(UpperCamelCase__  )
       transformers.utils.logging.enable_default_handler()
       transformers.utils.logging.enable_explicit_format()

       # Log on each process the small summary:
       logger.warning(
           f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
           + f"distributed training: {bool(training_args.local_rank != -1  )}, 16-bits training: {training_args.fpaa}"  )
       logger.info(f"Training/evaluation parameters {training_args}"  )

       # Detecting last checkpoint.
       A__											=						None
       if os.path.isdir(training_args.output_dir  ) and training_args.do_train and not training_args.overwrite_output_dir:
              A__											=						get_last_checkpoint(training_args.output_dir  )
              if last_checkpoint is None and len(os.listdir(training_args.output_dir  )  ) > 0:
                     raise ValueError(
                         f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                         '''Use --overwrite_output_dir to overcome.'''  )
              elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
                     logger.info(
                         f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                         '''the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'''  )

    # Initialize our dataset.
       A__											=						load_dataset(
           data_args.dataset_name ,     data_args.dataset_config_name ,     data_files=data_args.data_files ,     cache_dir=model_args.cache_dir ,     use_auth_token=True if model_args.use_auth_token else None ,     )

       # If we don't have a validation split, split off a percentage of train as validation.
       A__											=						None if '''validation''' in ds.keys() else data_args.train_val_split
       if isinstance(data_args.train_val_split ,     UpperCamelCase__  ) and data_args.train_val_split > 0.0:
              A__											=						ds['''train'''].train_test_split(data_args.train_val_split  )
              A__											=						split['''train''']
              A__											=						split['''test''']

       # Load pretrained model and image processor
       #
       # Distributed training:
       # The .from_pretrained methods guarantee that only one local process can concurrently
       # download model & vocab.
       A__											=						{
           '''cache_dir''': model_args.cache_dir,
           '''revision''': model_args.model_revision,
           '''use_auth_token''': True if model_args.use_auth_token else None,
       }
       if model_args.config_name:
              A__											=						ViTMAEConfig.from_pretrained(model_args.config_name ,     **UpperCamelCase__  )
       elif model_args.model_name_or_path:
              A__											=						ViTMAEConfig.from_pretrained(model_args.model_name_or_path ,     **UpperCamelCase__  )
       else:
              A__											=						ViTMAEConfig()
              logger.warning('''You are instantiating a new config instance from scratch.'''  )
              if model_args.config_overrides is not None:
                     logger.info(f"Overriding config: {model_args.config_overrides}"  )
                     config.update_from_string(model_args.config_overrides  )
                     logger.info(f"New config: {config}"  )

    # adapt config
       config.update(
           {
               '''mask_ratio''': model_args.mask_ratio,
               '''norm_pix_loss''': model_args.norm_pix_loss,
           }  )

       # create image processor
       if model_args.image_processor_name:
              A__											=						ViTImageProcessor.from_pretrained(model_args.image_processor_name ,     **UpperCamelCase__  )
       elif model_args.model_name_or_path:
              A__											=						ViTImageProcessor.from_pretrained(model_args.model_name_or_path ,     **UpperCamelCase__  )
       else:
              A__											=						ViTImageProcessor()

       # create model
       if model_args.model_name_or_path:
              A__											=						ViTMAEForPreTraining.from_pretrained(
                  model_args.model_name_or_path ,     from_tf=bool('''.ckpt''' in model_args.model_name_or_path  ) ,     config=UpperCamelCase__ ,     cache_dir=model_args.cache_dir ,     revision=model_args.model_revision ,     use_auth_token=True if model_args.use_auth_token else None ,     )
       else:
              logger.info('''Training new model from scratch'''  )
              A__											=						ViTMAEForPreTraining(UpperCamelCase__  )

       if training_args.do_train:
              A__											=						ds['''train'''].column_names
       else:
              A__											=						ds['''validation'''].column_names

       if data_args.image_column_name is not None:
              A__											=						data_args.image_column_name
       elif "image" in column_names:
              A__											=						'''image'''
       elif "img" in column_names:
              A__											=						'''img'''
       else:
              A__											=						column_names[0]

       # transformations as done in original MAE paper
       # source: https://github.com/facebookresearch/mae/blob/main/main_pretrain.py
       if "shortest_edge" in image_processor.size:
              A__											=						image_processor.size['''shortest_edge''']
       else:
              A__											=						(image_processor.size['''height'''], image_processor.size['''width'''])
       A__											=						Compose(
           [
               Lambda(lambda UpperCamelCase__  : img.convert('''RGB'''  ) if img.mode != "RGB" else img  ),
               RandomResizedCrop(UpperCamelCase__ ,     scale=(0.2, 1.0) ,     interpolation=InterpolationMode.BICUBIC  ),
               RandomHorizontalFlip(),
               ToTensor(),
               Normalize(mean=image_processor.image_mean ,     std=image_processor.image_std  ),
           ]  )

       def preprocess_images(UpperCamelCase__			:       List[str]  ):
              A__											=						[transforms(UpperCamelCase__  ) for image in examples[image_column_name]]
              return examples

       if training_args.do_train:
              if "train" not in ds:
                     raise ValueError('''--do_train requires a train dataset'''  )
              if data_args.max_train_samples is not None:
                     A__											=						ds['''train'''].shuffle(seed=training_args.seed  ).select(range(data_args.max_train_samples  )  )
              # Set the training transforms
              ds["train"].set_transform(UpperCamelCase__  )

       if training_args.do_eval:
              if "validation" not in ds:
                     raise ValueError('''--do_eval requires a validation dataset'''  )
              if data_args.max_eval_samples is not None:
                     A__											=						(
                         ds['''validation'''].shuffle(seed=training_args.seed  ).select(range(data_args.max_eval_samples  )  )
                     )
              # Set the validation transforms
              ds["validation"].set_transform(UpperCamelCase__  )

       # Compute absolute learning rate
       A__											=						(
           training_args.train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
       )
       if training_args.base_learning_rate is not None:
              A__											=						training_args.base_learning_rate * total_train_batch_size / 2_56

       # Initialize our trainer
       A__											=						Trainer(
           model=UpperCamelCase__ ,     args=UpperCamelCase__ ,     train_dataset=ds['''train'''] if training_args.do_train else None ,     eval_dataset=ds['''validation'''] if training_args.do_eval else None ,     tokenizer=UpperCamelCase__ ,     data_collator=UpperCamelCase__ ,     )

       # Training
       if training_args.do_train:
              A__											=						None
              if training_args.resume_from_checkpoint is not None:
                     A__											=						training_args.resume_from_checkpoint
              elif last_checkpoint is not None:
                     A__											=						last_checkpoint
              A__											=						trainer.train(resume_from_checkpoint=UpperCamelCase__  )
              trainer.save_model()
              trainer.log_metrics('''train''' ,     train_result.metrics  )
              trainer.save_metrics('''train''' ,     train_result.metrics  )
              trainer.save_state()

       # Evaluation
       if training_args.do_eval:
              A__											=						trainer.evaluate()
              trainer.log_metrics('''eval''' ,     UpperCamelCase__  )
              trainer.save_metrics('''eval''' ,     UpperCamelCase__  )

       # Write model card and (optionally) push to hub
       A__											=						{
           '''tasks''': '''masked-auto-encoding''',
           '''dataset''': data_args.dataset_name,
           '''tags''': ['''masked-auto-encoding'''],
       }
       if training_args.push_to_hub:
              trainer.push_to_hub(**UpperCamelCase__  )
       else:
              trainer.create_model_card(**UpperCamelCase__  )


def 	UpperCamelCase__(       UpperCamelCase__			:       Union[str, Any]  )->Optional[Any]:
       # For xla_spawn (TPUs)
       main()


if __name__ == "__main__":
       main()












import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
       import jax

       from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
       from transformers.models.bert.modeling_flax_bert import FlaxBertModel
       from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel




@require_flax
class     SCREAMING_SNAKE_CASE__ (						unittest.TestCase      ):



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["bert-base-cased", "bert-large-uncased"]:
               with self.subTest(__lowerCamelCase ):
                      A__											=						AutoConfig.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )

                      A__											=						FlaxAutoModel.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["roberta-base", "roberta-large"]:
               with self.subTest(__lowerCamelCase ):
                      A__											=						AutoConfig.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )

                      A__											=						FlaxAutoModel.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["bert-base-cased", "bert-large-uncased"]:
               A__											=						AutoTokenizer.from_pretrained(__lowerCamelCase )
               A__											=						FlaxBertModel.from_pretrained(__lowerCamelCase )
               A__											=						tokenizer('''Do you support jax jitted function?''',return_tensors=TensorType.JAX )

               @jax.jit
               def eval(**__lowerCamelCase ):
                      return model(**__lowerCamelCase )

               eval(**__lowerCamelCase ).block_until_ready()



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["roberta-base", "roberta-large"]:
               A__											=						AutoTokenizer.from_pretrained(__lowerCamelCase )
               A__											=						FlaxRobertaModel.from_pretrained(__lowerCamelCase )
               A__											=						tokenizer('''Do you support jax jitted function?''',return_tensors=TensorType.JAX )

               @jax.jit
               def eval(**__lowerCamelCase ):
                      return model(**__lowerCamelCase )

               eval(**__lowerCamelCase ).block_until_ready()



 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(
            __lowerCamelCase,'''bert-base is not a local folder and is not a valid model identifier''' ):
               A__											=						FlaxAutoModel.from_pretrained('''bert-base''' )



 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(
            __lowerCamelCase,r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)''' ):
               A__											=						FlaxAutoModel.from_pretrained(__lowerCamelCase,revision='''aaaaaa''' )



 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(
            __lowerCamelCase,'''hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack''',):
               A__											=						FlaxAutoModel.from_pretrained('''hf-internal-testing/config-no-model''' )




 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(__lowerCamelCase,'''Use `from_pt=True` to load this model''' ):
               A__											=						FlaxAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only''' )





