


"""simple docstring"""



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

import evaluate
import numpy as np
import torch
from datasets import load_dataset
from PIL import Image
from torchvision.transforms import (
    CenterCrop,
    Compose,
    Normalize,
    RandomHorizontalFlip,
    RandomResizedCrop,
    Resize,
    ToTensor,
)

import transformers
from transformers import (
    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    AutoModelForImageClassification,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
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
__snake_case       =       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-classification/requirements.txt""")

__snake_case       =       list(MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING.keys())
__snake_case       =       tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)







def     __lowerCAmelCase		(		lowercase    :    str					)      ->		Dict:


    """simple docstring"""





    with open(lowercase		,				"rb"					) as f:
        snake_case       : Union[str, Any]    =       Image.open(lowercase					)
        return im.convert("RGB"					)







@dataclass
class _lowerCAmelCase							:
   __UpperCAmelCase					:			Optional[str]				    =    field(
       default=snake_case_       ,				metadata={
           '''help''': '''Name of a dataset from the hub (could be your own, possibly private dataset hosted on the hub).'''
       }       ,				)
   __UpperCAmelCase					:			Optional[str]				    =    field(
       default=snake_case_       ,				metadata={'''help''': '''The configuration name of the dataset to use (via the datasets library).'''}      )
   __UpperCAmelCase					:			Optional[str]				    =    field(default=snake_case_       ,				metadata={'''help''': '''A folder containing the training data.'''}      )
   __UpperCAmelCase					:			Optional[str]				    =    field(default=snake_case_       ,				metadata={'''help''': '''A folder containing the validation data.'''}      )
   __UpperCAmelCase					:			Optional[float]				    =    field(
       default=0.15       ,				metadata={'''help''': '''Percent to split off of train for validation.'''}      )
   __UpperCAmelCase					:			Optional[int]				    =    field(
       default=snake_case_       ,				metadata={
           '''help''': (
               '''For debugging purposes or quicker training, truncate the number of training examples to this '''
               '''value if set.'''
           )
       }       ,				)
   __UpperCAmelCase					:			Optional[int]				    =    field(
       default=snake_case_       ,				metadata={
           '''help''': (
               '''For debugging purposes or quicker training, truncate the number of evaluation examples to this '''
               '''value if set.'''
           )
       }       ,				)







   def 					lowerCamelCase						(      self		)      ->						List[str]:

       '''simple docstring'''




       if self.dataset_name is None and (self.train_dir is None and self.validation_dir is None):
           raise ValueError(
               "You must specify either a dataset name from the hub or a train and/or validation directory."		)







@dataclass
class _lowerCAmelCase							:
   __UpperCAmelCase					:			str				    =    field(
       default='''google/vit-base-patch16-224-in21k'''       ,				metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}       ,				)
   __UpperCAmelCase					:			Optional[str]				    =    field(
       default=snake_case_       ,				metadata={'''help''': '''If training from scratch, pass a model type from the list: ''' + ''', '''.join(snake_case_      )}       ,				)
   __UpperCAmelCase					:			Optional[str]				    =    field(
       default=snake_case_       ,				metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''}      )
   __UpperCAmelCase					:			Optional[str]				    =    field(
       default=snake_case_       ,				metadata={'''help''': '''Where do you want to store the pretrained models downloaded from s3'''}      )
   __UpperCAmelCase					:			str				    =    field(
       default='''main'''       ,				metadata={'''help''': '''The specific model version to use (can be a branch name, tag name or commit id).'''}       ,				)
   __UpperCAmelCase					:			str				    =    field(default=snake_case_       ,				metadata={'''help''': '''Name or path of preprocessor config.'''}      )
   __UpperCAmelCase					:			bool				    =    field(
       default=snake_case_       ,				metadata={
           '''help''': (
               '''Will use the token generated when running `huggingface-cli login` (necessary to use this script '''
               '''with private models).'''
           )
       }       ,				)
   __UpperCAmelCase					:			bool				    =    field(
       default=snake_case_       ,				metadata={'''help''': '''Will enable to load a pretrained model whose head dimensions are different.'''}       ,				)







def     __lowerCAmelCase		(		lowercase    :    Any					)      ->		Union[str, Any]:


    """simple docstring"""





    snake_case       : Dict    =       torch.stack([example["pixel_values"] for example in examples]					)
    snake_case       : Any    =       torch.tensor([example["labels"] for example in examples]					)
    return {"pixel_values": pixel_values, "labels": labels}







def     __lowerCAmelCase		(		)      ->		str:


    """simple docstring"""





    snake_case       : Union[str, Any]    =       HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)					)
    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.
        snake_case   ,snake_case   ,snake_case       : List[Any]    =       parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]					)					)
    else:
        snake_case   ,snake_case   ,snake_case       : Tuple    =       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_image_classification"		,				lowercase		,				lowercase					)

    # 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()

    snake_case       : Dict    =       training_args.get_process_log_level()
    logger.setLevel(lowercase					)
    transformers.utils.logging.set_verbosity(lowercase					)
    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.
    snake_case       : List[str]    =       None
    if os.path.isdir(training_args.output_dir					) and training_args.do_train and not training_args.overwrite_output_dir:
        snake_case       : Any    =       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."					)

    # Set seed before initializing model.
    set_seed(training_args.seed					)

    # Initialize our dataset and prepare it for the 'image-classification' task.
    if data_args.dataset_name is not None:
        snake_case       : Tuple    =       load_dataset(
            data_args.dataset_name		,				data_args.dataset_config_name		,				cache_dir=model_args.cache_dir		,				task="image-classification"		,				use_auth_token=True if model_args.use_auth_token else None		,				)
    else:
        snake_case       : List[str]    =       {}
        if data_args.train_dir is not None:
            snake_case       : str    =       os.path.join(data_args.train_dir		,				"**"					)
        if data_args.validation_dir is not None:
            snake_case       : Optional[int]    =       os.path.join(data_args.validation_dir		,				"**"					)
        snake_case       : Optional[int]    =       load_dataset(
            "imagefolder"		,				data_files=lowercase		,				cache_dir=model_args.cache_dir		,				task="image-classification"		,				)

    # If we don't have a validation split, split off a percentage of train as validation.
    snake_case       : Optional[Any]    =       None if "validation" in dataset.keys() else data_args.train_val_split
    if isinstance(data_args.train_val_split		,				lowercase					) and data_args.train_val_split > 0.0:
        snake_case       : Optional[Any]    =       dataset["train"].train_test_split(data_args.train_val_split					)
        snake_case       : Tuple    =       split["train"]
        snake_case       : Optional[Any]    =       split["test"]

    # Prepare label mappings.
    # We'll include these in the model's config to get human readable labels in the Inference API.
    snake_case       : str    =       dataset["train"].features["labels"].names
    snake_case   ,snake_case       : List[str]    =       {}, {}
    for i, label in enumerate(lowercase					):
        snake_case       : str    =       str(lowercase					)
        snake_case       : Optional[int]    =       label

    # Load the accuracy metric from the datasets package
    snake_case       : Optional[int]    =       evaluate.load("accuracy"					)

    # Define our compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
    # predictions and label_ids field) and has to return a dictionary string to float.
    def compute_metrics(lowercase    :    List[str]					):
        return metric.compute(predictions=np.argmax(p.predictions		,				axis=1					)		,				references=p.label_ids					)

    snake_case       : List[Any]    =       AutoConfig.from_pretrained(
        model_args.config_name or model_args.model_name_or_path		,				num_labels=len(lowercase					)		,				labelaid=lowercase		,				idalabel=lowercase		,				finetuning_task="image-classification"		,				cache_dir=model_args.cache_dir		,				revision=model_args.model_revision		,				use_auth_token=True if model_args.use_auth_token else None		,				)
    snake_case       : Tuple    =       AutoModelForImageClassification.from_pretrained(
        model_args.model_name_or_path		,				from_tf=bool(".ckpt" in model_args.model_name_or_path					)		,				config=lowercase		,				cache_dir=model_args.cache_dir		,				revision=model_args.model_revision		,				use_auth_token=True if model_args.use_auth_token else None		,				ignore_mismatched_sizes=model_args.ignore_mismatched_sizes		,				)
    snake_case       : str    =       AutoImageProcessor.from_pretrained(
        model_args.image_processor_name or model_args.model_name_or_path		,				cache_dir=model_args.cache_dir		,				revision=model_args.model_revision		,				use_auth_token=True if model_args.use_auth_token else None		,				)

    # Define torchvision transforms to be applied to each image.
    if "shortest_edge" in image_processor.size:
        snake_case       : Optional[Any]    =       image_processor.size["shortest_edge"]
    else:
        snake_case       : Optional[int]    =       (image_processor.size["height"], image_processor.size["width"])
    snake_case       : Union[str, Any]    =       Normalize(mean=image_processor.image_mean		,				std=image_processor.image_std					)
    snake_case       : Any    =       Compose(
        [
            RandomResizedCrop(lowercase					),
            RandomHorizontalFlip(),
            ToTensor(),
            normalize,
        ]					)
    snake_case       : List[Any]    =       Compose(
        [
            Resize(lowercase					),
            CenterCrop(lowercase					),
            ToTensor(),
            normalize,
        ]					)

    def train_transforms(lowercase    :    Dict					):
        snake_case       : List[Any]    =       [
            _train_transforms(pil_img.convert("RGB"					)					) for pil_img in example_batch["image"]
        ]
        return example_batch

    def val_transforms(lowercase    :    str					):
        snake_case       : str    =       [_val_transforms(pil_img.convert("RGB"					)					) for pil_img in example_batch["image"]]
        return example_batch

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

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

    # Initalize our trainer
    snake_case       : List[Any]    =       Trainer(
        model=lowercase		,				args=lowercase		,				train_dataset=dataset["train"] if training_args.do_train else None		,				eval_dataset=dataset["validation"] if training_args.do_eval else None		,				compute_metrics=lowercase		,				tokenizer=lowercase		,				data_collator=lowercase		,				)

    # Training
    if training_args.do_train:
        snake_case       : List[Any]    =       None
        if training_args.resume_from_checkpoint is not None:
            snake_case       : Any    =       training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            snake_case       : Any    =       last_checkpoint
        snake_case       : Any    =       trainer.train(resume_from_checkpoint=lowercase					)
        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:
        snake_case       : Union[str, Any]    =       trainer.evaluate()
        trainer.log_metrics("eval"		,				lowercase					)
        trainer.save_metrics("eval"		,				lowercase					)

    # Write model card and (optionally) push to hub
    snake_case       : Dict    =       {
        "finetuned_from": model_args.model_name_or_path,
        "tasks": "image-classification",
        "dataset": data_args.dataset_name,
        "tags": ["image-classification", "vision"],
    }
    if training_args.push_to_hub:
        trainer.push_to_hub(**lowercase					)
    else:
        trainer.create_model_card(**lowercase					)


if __name__ == "__main__":
  main()









"""simple docstring"""



from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


__snake_case       =       logging.get_logger(__name__)

__snake_case       =       {
    """google/umt5-small""": """https://huggingface.co/google/umt5-small/resolve/main/config.json""",
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}







class _lowerCAmelCase							( snake_case_      ):
   __UpperCAmelCase					:			Dict				    =    '''umt5'''
   __UpperCAmelCase					:			str				    =    ['''past_key_values''']







   def __init__(      self		,      UpperCamelCase__=25_0112		,      UpperCamelCase__=512		,      UpperCamelCase__=64		,      UpperCamelCase__=1024		,      UpperCamelCase__=8		,      UpperCamelCase__=None		,      UpperCamelCase__=6		,      UpperCamelCase__=32		,      UpperCamelCase__=128		,      UpperCamelCase__=0.1		,      UpperCamelCase__=1e-6		,      UpperCamelCase__=1.0		,      UpperCamelCase__="gated-gelu"		,      UpperCamelCase__=True		,      UpperCamelCase__=True		,      UpperCamelCase__="T5Tokenizer"		,      UpperCamelCase__=True		,      UpperCamelCase__=0		,      UpperCamelCase__=1		,      UpperCamelCase__=0		,      **UpperCamelCase__		,      )      ->						Dict:

       '''simple docstring'''




       super().__init__(
           is_encoder_decoder=UpperCamelCase__		,      tokenizer_class=UpperCamelCase__		,      tie_word_embeddings=UpperCamelCase__		,      pad_token_id=UpperCamelCase__		,      eos_token_id=UpperCamelCase__		,      decoder_start_token_id=UpperCamelCase__		,      **UpperCamelCase__		,      )
       snake_case       : Optional[int]    =       vocab_size
       snake_case       : Any    =       d_model
       snake_case       : Optional[Any]    =       d_kv
       snake_case       : Any    =       d_ff
       snake_case       : Optional[Any]    =       num_layers
       snake_case       : Optional[Any]    =       (
           num_decoder_layers if num_decoder_layers is not None else self.num_layers
       )  # default = symmetry
       snake_case       : Optional[Any]    =       num_heads
       snake_case       : Optional[Any]    =       relative_attention_num_buckets
       snake_case       : Dict    =       relative_attention_max_distance
       snake_case       : Tuple    =       dropout_rate
       snake_case       : List[Any]    =       layer_norm_epsilon
       snake_case       : Any    =       initializer_factor
       snake_case       : Optional[Any]    =       feed_forward_proj
       snake_case       : Tuple    =       use_cache

       snake_case       : Any    =       self.feed_forward_proj.split("-"		)
       snake_case       : Tuple    =       act_info[-1]
       snake_case       : Optional[int]    =       act_info[0] == "gated"

       if len(UpperCamelCase__		) > 1 and act_info[0] != "gated" or len(UpperCamelCase__		) > 2:
           raise ValueError(
               F'`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.'
               "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
               "'gated-gelu' or 'relu'"		)

       if feed_forward_proj == "gated-gelu":
           snake_case       : List[str]    =       "gelu_new"







   @property
   def 					lowerCamelCase						(      self		)      ->						Any:

       '''simple docstring'''




       return self.d_model







   @property
   def 					lowerCamelCase						(      self		)      ->						str:

       '''simple docstring'''




       return self.num_heads







   @property
   def 					lowerCamelCase						(      self		)      ->						Union[str, Any]:

       '''simple docstring'''




       return self.num_layers







class _lowerCAmelCase							( snake_case_      ):







   @property
   # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
   def 					lowerCamelCase						(      self		)      ->						Mapping[str, Mapping[int, str]]:

       '''simple docstring'''




       snake_case       : Optional[int]    =       {
           "input_ids": {0: "batch", 1: "encoder_sequence"},
           "attention_mask": {0: "batch", 1: "encoder_sequence"},
       }
       if self.use_past:
           snake_case       : Tuple    =       "past_encoder_sequence + sequence"
           snake_case       : int    =       {0: "batch"}
           snake_case       : int    =       {0: "batch", 1: "past_decoder_sequence + sequence"}
       else:
           snake_case       : List[Any]    =       {0: "batch", 1: "decoder_sequence"}
           snake_case       : Any    =       {0: "batch", 1: "decoder_sequence"}

       if self.use_past:
           self.fill_with_past_key_values_(UpperCamelCase__		,      direction="inputs"		)

       return common_inputs







   @property
   # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
   def 					lowerCamelCase						(      self		)      ->						int:

       '''simple docstring'''




       return 13







   @property
   def 					lowerCamelCase						(      self		)      ->						float:

       '''simple docstring'''




       return 5e-4





