code: def main():
      # 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.
      parser = HfArgumentParser(
          (ModelArguments, DataTrainingArguments, TFTrainingArguments)
      )
      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.
          model_args, data_args, training_args = parser.parse_json_file(
              json_file=os.path.abspath(sys.argv[1])
          )
      else:
          (
              model_args,
              data_args,
              training_args,
          ) = parser.parse_args_into_dataclasses()
  
      if model_args.use_auth_token is not None:
          warnings.warn(
              "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead.",
              FutureWarning,
          )
          if model_args.token is not None:
              raise ValueError(
                  "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
              )
          model_args.token = model_args.use_auth_token
  
      if not (
          training_args.do_train
          or training_args.do_eval
          or training_args.do_predict
      ):
          exit(
              "Must specify at least one of --do_train, --do_eval or --do_predict!"
          )
  
      # 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/TensorFlow versions.
      send_example_telemetry(
          "run_image_classification",
          model_args,
          data_args,
          framework="tensorflow",
      )
  
      # Checkpoints. Find the checkpoint the use when loading the model.
      checkpoint = None
      if (
          os.path.isdir(training_args.output_dir)
          and training_args.do_train
          and not training_args.overwrite_output_dir
      ):
          checkpoint = get_last_checkpoint(training_args.output_dir)
          if (
              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 (
              checkpoint is not None
              and training_args.resume_from_checkpoint is None
          ):
              logger.info(
                  f"Checkpoint detected, resuming training at {checkpoint}. To avoid this behavior, change "
                  "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
              )
  
      # 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)],
      )
      log_level = training_args.get_process_log_level()
      logger.setLevel(log_level)
  
      # Set the verbosity to info of the Transformers logger (on main process only):
      if is_main_process(training_args.local_rank):
          transformers.utils.logging.set_verbosity_info()
          transformers.utils.logging.enable_default_handler()
          transformers.utils.logging.enable_explicit_format()
      logger.info(f"Training/evaluation parameters {training_args}")
  
      # region Dataset and labels
      # 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:
          dataset = load_dataset(
              data_args.dataset_name,
              data_args.dataset_config_name,
              cache_dir=model_args.cache_dir,
              task="image-classification",
              token=model_args.token,
          )
      else:
          data_files = {}
          if data_args.train_dir is not None:
              data_files["train"] = os.path.join(
                  data_args.train_dir, "**"
              )
          if data_args.validation_dir is not None:
              data_files["validation"] = os.path.join(
                  data_args.validation_dir, "**"
              )
          dataset = load_dataset(
              "imagefolder",
              data_files=data_files,
              cache_dir=model_args.cache_dir,
              task="image-classification",
          )
      # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
      # https://huggingface.co/docs/datasets/loading_datasets.
  
      # Prepare label mappings.
      # We'll include these in the model's config to get human readable labels in the Inference API.
      labels = dataset["train"].features["labels"].names
      label2id, id2label = {}, {}
      for i, label in enumerate(labels):
          label2id[label] = str(i)
          id2label[str(i)] = label
  
      # Load model image processor and configuration
      config = AutoConfig.from_pretrained(
          model_args.config_name or model_args.model_name_or_path,
          num_labels=len(labels),
          label2id=label2id,
          id2label=id2label,
          finetuning_task="image-classification",
          cache_dir=model_args.cache_dir,
          revision=model_args.model_revision,
          token=model_args.token,
          trust_remote_code=model_args.trust_remote_code,
      )
      image_processor = 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,
          token=model_args.token,
          trust_remote_code=model_args.trust_remote_code,
      )
  
      # If we don't have a validation split, split off a percentage of train as validation.
      data_args.train_val_split = (
          None
          if "validation" in dataset.keys()
          else data_args.train_val_split
      )
      if (
          isinstance(data_args.train_val_split, float)
          and data_args.train_val_split > 0.0
      ):
          split = dataset["train"].train_test_split(
              data_args.train_val_split
          )
          dataset["train"] = split["train"]
          dataset["validation"] = split["test"]
  
      # Define our data preprocessing function. It takes an image file path as input and returns
      # Write a note describing the resizing behaviour.
      if "shortest_edge" in image_processor.size:
          # We instead set the target size as (shortest_edge, shortest_edge) to here to ensure all images are batchable.
          image_size = (
              image_processor.size["shortest_edge"],
              image_processor.size["shortest_edge"],
          )
      else:
          image_size = (
              image_processor.size["height"],
              image_processor.size["width"],
          )
  
      def _train_transforms(image):
          img_size = image_size
          image = tf.keras.utils.img_to_array(image)
          image = random_resized_crop(image, size=img_size)
          image = tf.image.random_flip_left_right(image)
          image /= 255.0
          image = (
              image - image_processor.image_mean
          ) / image_processor.image_std
          image = tf.transpose(image, perm=[2, 0, 1])
          return image
  
      def _val_transforms(image):
          image = tf.keras.utils.img_to_array(image)
          image = tf.image.resize(image, size=image_size)
          # image = np.array(image) # FIXME - use tf.image function
          image = center_crop(image, size=image_size)
          image /= 255.0
          image = (
              image - image_processor.image_mean
          ) / image_processor.image_std
          image = tf.transpose(image, perm=[2, 0, 1])
          return image
  
      def train_transforms(example_batch):
          """Apply _train_transforms across a batch."""
          example_batch["pixel_values"] = [
              _train_transforms(pil_img.convert("RGB"))
              for pil_img in example_batch["image"]
          ]
          return example_batch
  
      def val_transforms(example_batch):
          """Apply _val_transforms across a batch."""
          example_batch["pixel_values"] = [
              _val_transforms(pil_img.convert("RGB"))
              for pil_img in example_batch["image"]
          ]
          return example_batch
  
      train_dataset = None
      if training_args.do_train:
          if "train" not in dataset:
              raise ValueError("--do_train requires a train dataset")
          train_dataset = dataset["train"]
          if data_args.max_train_samples is not None:
              train_dataset = train_dataset.shuffle(
                  seed=training_args.seed
              ).select(range(data_args.max_train_samples))
          train_dataset = train_dataset.map(
              train_transforms,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              load_from_cache_file=not data_args.overwrite_cache,
          )
  
      eval_dataset = None
      if training_args.do_eval:
          if "validation" not in dataset:
              raise ValueError(
                  "--do_eval requires a validation dataset"
              )
          eval_dataset = dataset["validation"]
          if data_args.max_eval_samples is not None:
              eval_dataset = eval_dataset.select(
                  range(data_args.max_eval_samples)
              )
          # Set the validation transforms
          eval_dataset = eval_dataset.map(
              val_transforms,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              load_from_cache_file=not data_args.overwrite_cache,
          )
  
      predict_dataset = None
      if training_args.do_predict:
          if "test" not in dataset:
              raise ValueError("--do_predict requires a test dataset")
          predict_dataset = dataset["test"]
          if data_args.max_predict_samples is not None:
              predict_dataset = predict_dataset.select(
                  range(data_args.max_predict_samples)
              )
          # Set the test transforms
          predict_dataset = predict_dataset.map(
              val_transforms,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              load_from_cache_file=not data_args.overwrite_cache,
          )
  
      collate_fn = DefaultDataCollator(return_tensors="np")
  
      # Load the accuracy metric from the datasets package
      metric = evaluate.load("accuracy", cache_dir=model_args.cache_dir)
  
      # 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(p):
          """Computes accuracy on a batch of predictions"""
          logits, label_ids = p
          predictions = np.argmax(logits, axis=-1)
          metrics = metric.compute(
              predictions=predictions, references=label_ids
          )
          return metrics
  
      with training_args.strategy.scope():
          if checkpoint is None:
              model_path = model_args.model_name_or_path
          else:
              model_path = checkpoint
  
          model = TFAutoModelForImageClassification.from_pretrained(
              model_path,
              config=config,
              from_pt=bool(".bin" in model_path),
              cache_dir=model_args.cache_dir,
              revision=model_args.model_revision,
              token=model_args.token,
              trust_remote_code=model_args.trust_remote_code,
              ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
          )
          num_replicas = training_args.strategy.num_replicas_in_sync
          total_train_batch_size = (
              training_args.per_device_train_batch_size * num_replicas
          )
          total_eval_batch_size = (
              training_args.per_device_eval_batch_size * num_replicas
          )
  
          dataset_options = tf.data.Options()
          dataset_options.experimental_distribute.auto_shard_policy = (
              tf.data.experimental.AutoShardPolicy.OFF
          )
  
          if training_args.do_train:
              num_train_steps = int(
                  len(train_dataset) * training_args.num_train_epochs
              )
              if training_args.warmup_steps > 0:
                  num_warmpup_steps = int(training_args.warmup_steps)
              elif training_args.warmup_ratio > 0:
                  num_warmpup_steps = int(
                      training_args.warmup_ratio * num_train_steps
                  )
              else:
                  num_warmpup_steps = 0
  
              optimizer, _ = create_optimizer(
                  init_lr=training_args.learning_rate,
                  num_train_steps=num_train_steps,
                  num_warmup_steps=num_warmpup_steps,
                  adam_beta1=training_args.adam_beta1,
                  adam_beta2=training_args.adam_beta2,
                  adam_epsilon=training_args.adam_epsilon,
                  weight_decay_rate=training_args.weight_decay,
                  adam_global_clipnorm=training_args.max_grad_norm,
              )
              # model.prepare_tf_dataset() wraps a Hugging Face dataset in a tf.data.Dataset which is ready to use in
              # training. This is the recommended way to use a Hugging Face dataset when training with Keras. You can also
              # use the lower-level dataset.to_tf_dataset() method, but you will have to specify things like column names
              # yourself if you use this method, whereas they are automatically inferred from the model input names when
              # using model.prepare_tf_dataset()
              # For more info see the docs:
              # https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.TFPreTrainedModel.prepare_tf_dataset
              # https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.to_tf_dataset
              train_dataset = model.prepare_tf_dataset(
                  train_dataset,
                  shuffle=True,
                  batch_size=total_train_batch_size,
                  collate_fn=collate_fn,
              ).with_options(dataset_options)
          else:
              optimizer = None
  
          if training_args.do_eval:
              eval_dataset = model.prepare_tf_dataset(
                  eval_dataset,
                  shuffle=False,
                  batch_size=total_eval_batch_size,
                  collate_fn=collate_fn,
              ).with_options(dataset_options)
  
          if training_args.do_predict:
              predict_dataset = model.prepare_tf_dataset(
                  predict_dataset,
                  shuffle=False,
                  batch_size=total_eval_batch_size,
                  collate_fn=collate_fn,
              ).with_options(dataset_options)
  
          # Transformers models compute the right loss for their task by default when labels are passed, and will
          # use this for training unless you specify your own loss function in compile().
          model.compile(
              optimizer=optimizer,
              jit_compile=training_args.xla,
              metrics=["accuracy"],
          )
  
          push_to_hub_model_id = training_args.push_to_hub_model_id
          if not push_to_hub_model_id:
              model_name = model_args.model_name_or_path.split("/")[-1]
              push_to_hub_model_id = (
                  f"{model_name}-finetuned-image-classification"
              )
  
          model_card_kwargs = {
              "finetuned_from": model_args.model_name_or_path,
              "tasks": "image-classification",
              "dataset": data_args.dataset_name,
              "tags": ["image-classification", "tensorflow", "vision"],
          }
  
          callbacks = []
          if eval_dataset is not None:
              callbacks.append(
                  KerasMetricCallback(
                      metric_fn=compute_metrics,
                      eval_dataset=eval_dataset,
                  )
              )
          if training_args.push_to_hub:
              callbacks.append(
                  PushToHubCallback(
                      output_dir=training_args.output_dir,
                      hub_model_id=push_to_hub_model_id,
                      hub_token=training_args.push_to_hub_token,
                      tokenizer=image_processor,
                      **model_card_kwargs,
                  )
              )
  
          if training_args.do_train:
              model.fit(
                  train_dataset,
                  validation_data=eval_dataset,
                  epochs=int(training_args.num_train_epochs),
                  callbacks=callbacks,
              )
  
          if training_args.do_eval:
              n_eval_batches = len(eval_dataset)
              eval_predictions = model.predict(
                  eval_dataset, steps=n_eval_batches
              )
              eval_labels = dataset["validation"]["labels"][
                  : n_eval_batches * total_eval_batch_size
              ]
              eval_metrics = compute_metrics(
                  (eval_predictions.logits, eval_labels)
              )
              logging.info("Eval metrics:")
              for metric_name, value in eval_metrics.items():
                  logging.info(f"{metric_name}: {value:.3f}")
  
          if training_args.output_dir is not None:
              os.makedirs(training_args.output_dir, exist_ok=True)
              with open(
                  os.path.join(
                      training_args.output_dir, "all_results.json"
                  ),
                  "w",
              ) as f:
                  f.write(json.dumps(eval_metrics))
  
          if training_args.do_predict:
              n_predict_batches = len(predict_dataset)
              test_predictions = model.predict(
                  predict_dataset, steps=n_predict_batches
              )
              test_labels = dataset["validation"]["labels"][
                  : n_predict_batches * total_eval_batch_size
              ]
              test_metrics = compute_metrics(
                  (test_predictions.logits, test_labels)
              )
              logging.info("Test metrics:")
              for metric_name, value in test_metrics.items():
                  logging.info(f"{metric_name}: {value:.3f}")
  
          if (
              training_args.output_dir is not None
              and not training_args.push_to_hub
          ):
              # If we're not pushing to hub, at least save a local copy when we're done
              model.save_pretrained(training_args.output_dir)
  
imports: import json
  import logging
  import os
  import sys
  import warnings
  from dataclasses import dataclass, field
  from typing import Optional
  import evaluate
  import numpy
  import tensorflow
  from datasets import load_dataset
  from PIL import Image
  import transformers
  from transformers import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING, AutoConfig, AutoImageProcessor, DefaultDataCollator, HfArgumentParser, PushToHubCallback, TFAutoModelForImageClassification, TFTrainingArguments, create_optimizer, set_seed
  from transformers.keras_callbacks import KerasMetricCallback
  from transformers.trainer_utils import get_last_checkpoint, is_main_process
  from transformers.utils import check_min_version, send_example_telemetry
  from transformers.utils.versions import require_version
label_desc: Use image node_id 5 for calling a global function with example usage: main() and returns: image, image, example_batch, example_batch, metrics