




"""simple docstring"""



import argparse
import logging
import os
import sys
import time

import tensorflow as tf
from datasets import load_dataset
from tqdm import tqdm

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
from transformers.utils import is_sagemaker_dp_enabled


if os.environ.get("SDP_ENABLED") or is_sagemaker_dp_enabled():
      SCREAMING_SNAKE_CASE							=			True
      SCREAMING_SNAKE_CASE							=			"p3dn.24xlarge"
      import smdistributed.dataparallel.tensorflow as sdp
else:
      SCREAMING_SNAKE_CASE							=			False







def     _SCREAMING_SNAKE_CASE  (				lowercase_     ,					lowercase_     ,					lowercase_     ,					lowercase_     ,					lowercase_     ,					lowercase_     ,					lowercase_=None		)		->						Union[str, Any]:
 A__			=							tqdm(lowercase_		)
 for i, batch in enumerate(lowercase_		):
  with tf.GradientTape() as tape:
   A__, A__			=							batch
   A__			=							model(lowercase_		)
   A__			=							loss(lowercase_     ,					outputs.logits		)

  if SDP_ENABLED:
   A__			=							sdp.DistributedGradientTape(lowercase_     ,					sparse_as_dense=lowercase_		)

  A__			=							tape.gradient(lowercase_     ,					model.trainable_variables		)
  opt.apply_gradients(zip(lowercase_     ,					model.trainable_variables		)		)

  pbar.set_description(f"""Loss: {loss_value:.4f}"""		)

  if SDP_ENABLED and i == 0:
   sdp.broadcast_variables(model.variables     ,					root_rank=0		)
   sdp.broadcast_variables(opt.variables()     ,					root_rank=0		)

  if max_steps and i >= max_steps:
   break

 A__			=							{"loss": loss_value.numpy()}
 return train_results







def     _SCREAMING_SNAKE_CASE  (				lowercase_     ,					lowercase_     ,					lowercase_		)		->						List[Any]:
 # Load dataset
 A__, A__			=							load_dataset("imdb"     ,					split=["train", "test"]		)

 # Preprocess train dataset
 A__			=							train_dataset.map(
     lambda lowercase_		: tokenizer(e["text"]     ,					truncation=lowercase_     ,					padding="max_length"		)     ,					batched=lowercase_		)
 train_dataset.set_format(type="tensorflow"     ,					columns=["input_ids", "attention_mask", "label"]		)

 A__			=							{
     x: train_dataset[x].to_tensor(default_value=0     ,					shape=[None, tokenizer.model_max_length]		)
     for x in ["input_ids", "attention_mask"]
 }
 A__			=							tf.data.Dataset.from_tensor_slices((train_features, train_dataset["label"])		)

 # Preprocess test dataset
 A__			=							test_dataset.map(
     lambda lowercase_		: tokenizer(e["text"]     ,					truncation=lowercase_     ,					padding="max_length"		)     ,					batched=lowercase_		)
 test_dataset.set_format(type="tensorflow"     ,					columns=["input_ids", "attention_mask", "label"]		)

 A__			=							{
     x: test_dataset[x].to_tensor(default_value=0     ,					shape=[None, tokenizer.model_max_length]		)
     for x in ["input_ids", "attention_mask"]
 }
 A__			=							tf.data.Dataset.from_tensor_slices((test_features, test_dataset["label"])		)

 if SDP_ENABLED:
  A__			=							tf_train_dataset.shard(sdp.size()     ,					sdp.rank()		)
  A__			=							tf_test_dataset.shard(sdp.size()     ,					sdp.rank()		)
 A__			=							tf_train_dataset.batch(lowercase_     ,					drop_remainder=lowercase_		)
 A__			=							tf_test_dataset.batch(lowercase_     ,					drop_remainder=lowercase_		)

 return tf_train_dataset, tf_test_dataset


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE							=			argparse.ArgumentParser()

      # Hyperparameters sent by the client are passed as command-line arguments to the script.
      parser.add_argument("--epochs", type=int, default=3)
      parser.add_argument("--per_device_train_batch_size", type=int, default=16)
      parser.add_argument("--per_device_eval_batch_size", type=int, default=8)
      parser.add_argument("--model_name_or_path", type=str)
      parser.add_argument("--learning_rate", type=str, default=5E-5)
      parser.add_argument("--do_train", type=bool, default=True)
      parser.add_argument("--do_eval", type=bool, default=True)
      parser.add_argument("--output_dir", type=str)
      parser.add_argument("--max_steps", type=int, default=None)

      # Data, model, and output directories
      parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
      parser.add_argument("--model_dir", type=str, default=os.environ["SM_MODEL_DIR"])
      parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])

      SCREAMING_SNAKE_CASE,   SCREAMING_SNAKE_CASE							=			parser.parse_known_args()

      # Set up logging
      SCREAMING_SNAKE_CASE							=			logging.getLogger(__name__)

      logging.basicConfig(
          level=logging.getLevelName("INFO"),
          handlers=[logging.StreamHandler(sys.stdout)],
          format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
      )

      if SDP_ENABLED:
            sdp.init()

            SCREAMING_SNAKE_CASE							=			tf.config.experimental.list_physical_devices("GPU")
            for gpu in gpus:
                  tf.config.experimental.set_memory_growth(gpu, True)
            if gpus:
                  tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], "GPU")

    # Load model and tokenizer
      SCREAMING_SNAKE_CASE							=			TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
      SCREAMING_SNAKE_CASE							=			AutoTokenizer.from_pretrained(args.model_name_or_path)

      # get datasets
      SCREAMING_SNAKE_CASE,   SCREAMING_SNAKE_CASE							=			get_datasets(
          tokenizer=tokenizer,
          train_batch_size=args.per_device_train_batch_size,
          eval_batch_size=args.per_device_eval_batch_size,
      )

      # fine optimizer and loss
      SCREAMING_SNAKE_CASE							=			tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
      SCREAMING_SNAKE_CASE							=			tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
      SCREAMING_SNAKE_CASE							=			[tf.keras.metrics.SparseCategoricalAccuracy()]
      model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

      # Training
      if args.do_train:
            # train_results = model.fit(tf_train_dataset, epochs=args.epochs, batch_size=args.train_batch_size)
            SCREAMING_SNAKE_CASE							=			time.time()
            SCREAMING_SNAKE_CASE							=			fit(
                model,
                loss,
                optimizer,
                tf_train_dataset,
                args.epochs,
                args.per_device_train_batch_size,
                max_steps=args.max_steps,
            )
            SCREAMING_SNAKE_CASE							=			time.time() - start_train_time
            logger.info("*** Train ***")
            logger.info(f'train_runtime = {end_train_time}')

            SCREAMING_SNAKE_CASE							=			os.path.join(args.output_dir, "train_results.txt")

            if not SDP_ENABLED or sdp.rank() == 0:
                  with open(output_eval_file, "w") as writer:
                        logger.info("***** Train results *****")
                        logger.info(train_results)
                        for key, value in train_results.items():
                              logger.info(f'  {key} = {value}')
                              writer.write(f'{key} = {value}\n')

    # Evaluation
      if args.do_eval and (not SDP_ENABLED or sdp.rank() == 0):
            SCREAMING_SNAKE_CASE							=			model.evaluate(tf_test_dataset, batch_size=args.per_device_eval_batch_size, return_dict=True)
            logger.info("*** Evaluate ***")

            SCREAMING_SNAKE_CASE							=			os.path.join(args.output_dir, "eval_results.txt")

            with open(output_eval_file, "w") as writer:
                  logger.info("***** Eval results *****")
                  logger.info(result)
                  for key, value in result.items():
                        logger.info(f'  {key} = {value}')
                        writer.write(f'{key} = {value}\n')

    # Save result
      if SDP_ENABLED:
            if sdp.rank() == 0:
                  model.save_pretrained(args.output_dir)
                  tokenizer.save_pretrained(args.output_dir)
      else:
            model.save_pretrained(args.output_dir)
            tokenizer.save_pretrained(args.output_dir)











"""simple docstring"""



import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import ViTImageProcessor




class       UpperCAmelCase_    (		unittest.TestCase				):




    def __init__(       self		:			Union[str, Any]      ,					snake_case_		:			Tuple      ,					snake_case_		:			Dict=7      ,					snake_case_		:			List[Any]=3      ,					snake_case_		:			Optional[int]=18      ,					snake_case_		:			Tuple=30      ,					snake_case_		:			Dict=400      ,					snake_case_		:			List[str]=True      ,					snake_case_		:			Optional[int]=None      ,					snake_case_		:			List[str]=True      ,					snake_case_		:			Optional[Any]=[0.5, 0.5, 0.5]      ,					snake_case_		:			Optional[int]=[0.5, 0.5, 0.5]      ,					) ->				List[Any]:




     '''simple docstring'''

     A__			=							size if size is not None else {"height": 18, "width": 18}
     A__			=							parent
     A__			=							batch_size
     A__			=							num_channels
     A__			=							image_size
     A__			=							min_resolution
     A__			=							max_resolution
     A__			=							do_resize
     A__			=							size
     A__			=							do_normalize
     A__			=							image_mean
     A__			=							image_std




    def 				__magic_name__				(       self		:			Union[str, Any] ) ->				Dict:




     '''simple docstring'''

     return {
         "image_mean": self.image_mean,
         "image_std": self.image_std,
         "do_normalize": self.do_normalize,
         "do_resize": self.do_resize,
         "size": self.size,
     }






@require_torch
@require_vision
class       UpperCAmelCase_    (		A_,  unittest.TestCase				):
    lowercase__		   = ViTImageProcessor if is_vision_available() else None




    def 				__magic_name__				(       self		:			int ) ->				Any:




     '''simple docstring'''

     A__			=							ViTImageProcessingTester(self )




    @property
    def 				__magic_name__				(       self		:			Optional[Any] ) ->				str:




     '''simple docstring'''

     return self.image_processor_tester.prepare_image_processor_dict()




    def 				__magic_name__				(       self		:			Tuple ) ->				str:




     '''simple docstring'''

     A__			=							self.image_processing_class(**self.image_processor_dict )
     self.assertTrue(hasattr(snake_case_      ,					"image_mean" ) )
     self.assertTrue(hasattr(snake_case_      ,					"image_std" ) )
     self.assertTrue(hasattr(snake_case_      ,					"do_normalize" ) )
     self.assertTrue(hasattr(snake_case_      ,					"do_resize" ) )
     self.assertTrue(hasattr(snake_case_      ,					"size" ) )




    def 				__magic_name__				(       self		:			Union[str, Any] ) ->				List[str]:




     '''simple docstring'''

     A__			=							self.image_processing_class.from_dict(self.image_processor_dict )
     self.assertEqual(image_processor.size      ,					{"height": 18, "width": 18} )

     A__			=							self.image_processing_class.from_dict(self.image_processor_dict      ,					size=42 )
     self.assertEqual(image_processor.size      ,					{"height": 42, "width": 42} )




    def 				__magic_name__				(       self		:			Any ) ->				Optional[Any]:




     '''simple docstring'''

     pass




    def 				__magic_name__				(       self		:			Optional[int] ) ->				Any:




     '''simple docstring'''

     A__			=							self.image_processing_class(**self.image_processor_dict )
     # create random PIL images
     A__			=							prepare_image_inputs(self.image_processor_tester      ,					equal_resolution=snake_case_ )
     for image in image_inputs:
      self.assertIsInstance(snake_case_      ,					Image.Image )

     # Test not batched input
     A__			=							image_processing(image_inputs[0]      ,					return_tensors="pt" ).pixel_values
     self.assertEqual(
         encoded_images.shape      ,					(
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size["height"],
             self.image_processor_tester.size["width"],
         )      ,					)

     # Test batched
     A__			=							image_processing(snake_case_      ,					return_tensors="pt" ).pixel_values
     self.assertEqual(
         encoded_images.shape      ,					(
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size["height"],
             self.image_processor_tester.size["width"],
         )      ,					)




    def 				__magic_name__				(       self		:			str ) ->				Optional[Any]:




     '''simple docstring'''

     A__			=							self.image_processing_class(**self.image_processor_dict )
     # create random numpy tensors
     A__			=							prepare_image_inputs(self.image_processor_tester      ,					equal_resolution=snake_case_      ,					numpify=snake_case_ )
     for image in image_inputs:
      self.assertIsInstance(snake_case_      ,					np.ndarray )

     # Test not batched input
     A__			=							image_processing(image_inputs[0]      ,					return_tensors="pt" ).pixel_values
     self.assertEqual(
         encoded_images.shape      ,					(
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size["height"],
             self.image_processor_tester.size["width"],
         )      ,					)

     # Test batched
     A__			=							image_processing(snake_case_      ,					return_tensors="pt" ).pixel_values
     self.assertEqual(
         encoded_images.shape      ,					(
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size["height"],
             self.image_processor_tester.size["width"],
         )      ,					)




    def 				__magic_name__				(       self		:			Dict ) ->				Union[str, Any]:




     '''simple docstring'''

     A__			=							self.image_processing_class(**self.image_processor_dict )
     # create random PyTorch tensors
     A__			=							prepare_image_inputs(self.image_processor_tester      ,					equal_resolution=snake_case_      ,					torchify=snake_case_ )
     for image in image_inputs:
      self.assertIsInstance(snake_case_      ,					torch.Tensor )

     # Test not batched input
     A__			=							image_processing(image_inputs[0]      ,					return_tensors="pt" ).pixel_values
     self.assertEqual(
         encoded_images.shape      ,					(
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size["height"],
             self.image_processor_tester.size["width"],
         )      ,					)

     # Test batched
     A__			=							image_processing(snake_case_      ,					return_tensors="pt" ).pixel_values
     self.assertEqual(
         encoded_images.shape      ,					(
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size["height"],
             self.image_processor_tester.size["width"],
         )      ,					)





