




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
_lowerCAmelCase  :      Union[str, Any]						=      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")

_lowerCAmelCase  :      List[str]						=      list(MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING.keys())
_lowerCAmelCase  :      Dict						=      tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)



def   lowerCAmelCase						(						_lowerCAmelCase							:      Optional[Any]		):




 """simple docstring"""







 with open(_lowerCAmelCase  ,							"rb"		) as f:
  UpperCAmelCase__							       =  Image.open(_lowerCAmelCase		)
  return im.convert("RGB"		)
@dataclass
class 	_UpperCamelCase					:
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={
              """help""": """Name of a dataset from the hub (could be your own, possibly private dataset hosted on the hub)."""
          }		,				)
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={"""help""": """The configuration name of the dataset to use (via the datasets library)."""}						)
      UpperCAmelCase_						=						field(default=UpperCamelCase__		,				metadata={"""help""": """A folder containing the training data."""}						)
      UpperCAmelCase_						=						field(default=UpperCamelCase__		,				metadata={"""help""": """A folder containing the validation data."""}						)
      UpperCAmelCase_						=						field(
          default=0.15		,				metadata={"""help""": """Percent to split off of train for validation."""}						)
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={
              """help""": (
                  """For debugging purposes or quicker training, truncate the number of training examples to this """
                  """value if set."""
              )
          }		,				)
      UpperCAmelCase_						=						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      :Any					)   ->	int:







       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 	_UpperCamelCase					:
      UpperCAmelCase_						=						field(
          default="""google/vit-base-patch16-224-in21k"""		,				metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}		,				)
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={"""help""": """If training from scratch, pass a model type from the list: """ + """, """.join(UpperCamelCase__						)}		,				)
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={"""help""": """Pretrained config name or path if not the same as model_name"""}						)
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={"""help""": """Where do you want to store the pretrained models downloaded from s3"""}						)
      UpperCAmelCase_						=						field(
          default="""main"""		,				metadata={"""help""": """The specific model version to use (can be a branch name, tag name or commit id)."""}		,				)
      UpperCAmelCase_						=						field(default=UpperCamelCase__		,				metadata={"""help""": """Name or path of preprocessor config."""}						)
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={
              """help""": (
                  """Will use the token generated when running `huggingface-cli login` (necessary to use this script """
                  """with private models)."""
              )
          }		,				)
      UpperCAmelCase_						=						field(
          default=UpperCamelCase__		,				metadata={"""help""": """Will enable to load a pretrained model whose head dimensions are different."""}		,				)



def   lowerCAmelCase						(						_lowerCAmelCase							:      List[Any]		):




 """simple docstring"""







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



def   lowerCAmelCase						(						):




 """simple docstring"""







 UpperCAmelCase__							       =  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.
  UpperCAmelCase__       ,       UpperCAmelCase__       ,       UpperCAmelCase__							       =  parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]		)		)
 else:
  UpperCAmelCase__       ,       UpperCAmelCase__       ,       UpperCAmelCase__							       =  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"  ,							_lowerCAmelCase  ,							_lowerCAmelCase		)

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

 UpperCAmelCase__							       =  training_args.get_process_log_level()
 logger.setLevel(_lowerCAmelCase		)
 transformers.utils.logging.set_verbosity(_lowerCAmelCase		)
 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.
 UpperCAmelCase__							       =  None
 if os.path.isdir(training_args.output_dir		) and training_args.do_train and not training_args.overwrite_output_dir:
  UpperCAmelCase__							       =  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:
  UpperCAmelCase__							       =  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:
  UpperCAmelCase__							       =  {}
  if data_args.train_dir is not None:
   UpperCAmelCase__							       =  os.path.join(data_args.train_dir  ,							"**"		)
  if data_args.validation_dir is not None:
   UpperCAmelCase__							       =  os.path.join(data_args.validation_dir  ,							"**"		)
  UpperCAmelCase__							       =  load_dataset(
      "imagefolder"  ,							data_files=_lowerCAmelCase  ,							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.
 UpperCAmelCase__							       =  None if "validation" in dataset.keys() else data_args.train_val_split
 if isinstance(data_args.train_val_split  ,							_lowerCAmelCase		) and data_args.train_val_split > 0.0:
  UpperCAmelCase__							       =  dataset["train"].train_test_split(data_args.train_val_split		)
  UpperCAmelCase__							       =  split["train"]
  UpperCAmelCase__							       =  split["test"]

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

 # Load the accuracy metric from the datasets package
 UpperCAmelCase__							       =  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(_lowerCAmelCase							:      str		):
  return metric.compute(predictions=np.argmax(p.predictions  ,							axis=1		)  ,							references=p.label_ids		)

 UpperCAmelCase__							       =  AutoConfig.from_pretrained(
     model_args.config_name or model_args.model_name_or_path  ,							num_labels=len(_lowerCAmelCase		)  ,							labelaid=_lowerCAmelCase  ,							idalabel=_lowerCAmelCase  ,							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  ,							)
 UpperCAmelCase__							       =  AutoModelForImageClassification.from_pretrained(
     model_args.model_name_or_path  ,							from_tf=bool(".ckpt" in model_args.model_name_or_path		)  ,							config=_lowerCAmelCase  ,							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  ,							)
 UpperCAmelCase__							       =  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:
  UpperCAmelCase__							       =  image_processor.size["shortest_edge"]
 else:
  UpperCAmelCase__							       =  (image_processor.size["height"], image_processor.size["width"])
 UpperCAmelCase__							       =  Normalize(mean=image_processor.image_mean  ,							std=image_processor.image_std		)
 UpperCAmelCase__							       =  Compose(
     [
         RandomResizedCrop(_lowerCAmelCase		),
         RandomHorizontalFlip(),
         ToTensor(),
         normalize,
     ]		)
 UpperCAmelCase__							       =  Compose(
     [
         Resize(_lowerCAmelCase		),
         CenterCrop(_lowerCAmelCase		),
         ToTensor(),
         normalize,
     ]		)

 def train_transforms(_lowerCAmelCase							:      Tuple		):
  UpperCAmelCase__							       =  [
      _train_transforms(pil_img.convert("RGB"		)		) for pil_img in example_batch["image"]
  ]
  return example_batch

 def val_transforms(_lowerCAmelCase							:      List[str]		):
  UpperCAmelCase__							       =  [_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:
   UpperCAmelCase__							       =  (
       dataset["train"].shuffle(seed=training_args.seed		).select(range(data_args.max_train_samples		)		)
   )
  # Set the training transforms
  dataset["train"].set_transform(_lowerCAmelCase		)

 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:
   UpperCAmelCase__							       =  (
       dataset["validation"].shuffle(seed=training_args.seed		).select(range(data_args.max_eval_samples		)		)
   )
  # Set the validation transforms
  dataset["validation"].set_transform(_lowerCAmelCase		)

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

 # Training
 if training_args.do_train:
  UpperCAmelCase__							       =  None
  if training_args.resume_from_checkpoint is not None:
   UpperCAmelCase__							       =  training_args.resume_from_checkpoint
  elif last_checkpoint is not None:
   UpperCAmelCase__							       =  last_checkpoint
  UpperCAmelCase__							       =  trainer.train(resume_from_checkpoint=_lowerCAmelCase		)
  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:
  UpperCAmelCase__							       =  trainer.evaluate()
  trainer.log_metrics("eval"  ,							_lowerCAmelCase		)
  trainer.save_metrics("eval"  ,							_lowerCAmelCase		)

 # Write model card and (optionally) push to hub
 UpperCAmelCase__							       =  {
     "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(**_lowerCAmelCase		)
 else:
  trainer.create_model_card(**_lowerCAmelCase		)


if __name__ == "__main__":
  main()









from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


lowerCAmelCase       =							logging.get_logger(__name__)
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				with open(SCREAMING_SNAKE_CASE						,				'''rb'''     ) as flax_state_f:
						lowercase__            =      from_bytes(SCREAMING_SNAKE_CASE						,				flax_state_f.read()     )
		except UnpicklingError as e:
				try:
						with open(SCREAMING_SNAKE_CASE     ) as f:
								if f.read().startswith('''version'''     ):
										raise OSError(
										    '''You seem to have cloned a repository without having git-lfs installed. Please'''
										    ''' install git-lfs and run `git lfs install` followed by `git lfs pull` in the'''
										    ''' folder you cloned.'''     )
								else:
										raise ValueError from e
				except (UnicodeDecodeError, ValueError):
						raise EnvironmentError(f'Unable to convert {model_file} to Flax deserializable object. '     )

		return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     )
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				import torch  # noqa: F401
		except ImportError:
				logger.error(
				    '''Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see'''
				    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
				    ''' instructions.'''     )
				raise

		# check if we have bf16 weights
		lowercase__            =      flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE     : x.dtype == jnp.bfloataa						,				SCREAMING_SNAKE_CASE     )     ).values()
		if any(SCREAMING_SNAKE_CASE     ):
				# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

				# and bf16 is not fully supported in PT yet.
				logger.warning(
				    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
				    '''before loading those in PyTorch model.'''     )
				lowercase__            =      jax.tree_util.tree_map(
				    lambda SCREAMING_SNAKE_CASE     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params						,				SCREAMING_SNAKE_CASE     )

		lowercase__            =      ''''''

		lowercase__            =      flatten_dict(SCREAMING_SNAKE_CASE						,				sep='''.'''     )
		lowercase__            =      pt_model.state_dict()

		# keep track of unexpected & missing keys
		lowercase__            =      []
		lowercase__            =      set(pt_model_dict.keys()     )

		for flax_key_tuple, flax_tensor in flax_state_dict.items():
				lowercase__            =      flax_key_tuple.split('''.'''     )

				if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      jnp.transpose(SCREAMING_SNAKE_CASE						,				(3, 2, 0, 1)     )
				elif flax_key_tuple_array[-1] == "kernel":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      flax_tensor.T
				elif flax_key_tuple_array[-1] == "scale":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']

				if "time_embedding" not in flax_key_tuple_array:
						for i, flax_key_tuple_string in enumerate(SCREAMING_SNAKE_CASE     ):
								lowercase__            =      (
								    flax_key_tuple_string.replace('''_0'''						,				'''.0'''     )
								    .replace('''_1'''						,				'''.1'''     )
								    .replace('''_2'''						,				'''.2'''     )
								    .replace('''_3'''						,				'''.3'''     )
								    .replace('''_4'''						,				'''.4'''     )
								    .replace('''_5'''						,				'''.5'''     )
								    .replace('''_6'''						,				'''.6'''     )
								    .replace('''_7'''						,				'''.7'''     )
								    .replace('''_8'''						,				'''.8'''     )
								    .replace('''_9'''						,				'''.9'''     )
								)

				lowercase__            =      '''.'''.join(SCREAMING_SNAKE_CASE     )

				if flax_key in pt_model_dict:
						if flax_tensor.shape != pt_model_dict[flax_key].shape:
								raise ValueError(
								    f'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
								    f'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'     )
						else:
								# add weight to pytorch dict
								lowercase__            =      np.asarray(SCREAMING_SNAKE_CASE     ) if not isinstance(SCREAMING_SNAKE_CASE						,				np.ndarray     ) else flax_tensor
								lowercase__            =      torch.from_numpy(SCREAMING_SNAKE_CASE     )
								# remove from missing keys
								missing_keys.remove(SCREAMING_SNAKE_CASE     )
				else:
						# weight is not expected by PyTorch model
						unexpected_keys.append(SCREAMING_SNAKE_CASE     )

		pt_model.load_state_dict(SCREAMING_SNAKE_CASE     )

		# re-transform missing_keys to list
		lowercase__            =      list(SCREAMING_SNAKE_CASE     )

		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
				    f' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
				    f' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
				    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
				    f' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
				    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
				    ''' FlaxBertForSequenceClassification model).'''     )
		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    f'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
				    f' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
				    ''' use it for predictions and inference.'''     )

		return pt_model


