




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_						:					Any			   =       True
      SCREAMING_SNAKE_CASE_						:					List[str]			   =       'p3dn.24xlarge'
      import smdistributed.dataparallel.tensorflow as sdp
else:
      SCREAMING_SNAKE_CASE_						:					Any			   =       False
def        __lowerCAmelCase (UpperCAmelCase_   ,    UpperCAmelCase_   ,    UpperCAmelCase_   ,    UpperCAmelCase_   ,    UpperCAmelCase_   ,    UpperCAmelCase_   ,    UpperCAmelCase_=None		):
    a__      =       tqdm(UpperCAmelCase_		)
    for i, batch in enumerate(UpperCAmelCase_		):
        with tf.GradientTape() as tape:
            a__ ,		a__      =       batch
            a__      =       model(UpperCAmelCase_		)
            a__      =       loss(UpperCAmelCase_   ,    outputs.logits		)

        if SDP_ENABLED:
            a__      =       sdp.DistributedGradientTape(UpperCAmelCase_   ,    sparse_as_dense=UpperCAmelCase_		)

        a__      =       tape.gradient(UpperCAmelCase_   ,    model.trainable_variables		)
        opt.apply_gradients(zip(UpperCAmelCase_   ,    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        __lowerCAmelCase (UpperCAmelCase_   ,    UpperCAmelCase_   ,    UpperCAmelCase_		):
    a__ ,		a__      =       load_dataset("""imdb"""   ,    split=["""train""", """test"""]		)

    # Preprocess train dataset
    a__      =       train_dataset.map(
        lambda UpperCAmelCase_		: tokenizer(e["""text"""]   ,    truncation=UpperCAmelCase_   ,    padding="""max_length"""		)   ,    batched=UpperCAmelCase_		)
    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 UpperCAmelCase_		: tokenizer(e["""text"""]   ,    truncation=UpperCAmelCase_   ,    padding="""max_length"""		)   ,    batched=UpperCAmelCase_		)
    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(UpperCAmelCase_   ,    drop_remainder=UpperCAmelCase_		)
    a__      =       tf_test_dataset.batch(UpperCAmelCase_   ,    drop_remainder=UpperCAmelCase_		)

    return tf_train_dataset, tf_test_dataset


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE_						:					Optional[int]			   =       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_						:					Union[str, Any]			   =       parser.parse_known_args()

      # Set up logging
      SCREAMING_SNAKE_CASE_						:					List[str]			   =       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_						:					Optional[Any]			   =       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_						:					Union[str, Any]			   =       TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
      SCREAMING_SNAKE_CASE_						:					Optional[int]			   =       AutoTokenizer.from_pretrained(args.model_name_or_path)

      # get datasets
      SCREAMING_SNAKE_CASE_						:					Optional[int]			   =       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_						:					List[str]			   =       tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
      SCREAMING_SNAKE_CASE_						:					Any			   =       tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
      SCREAMING_SNAKE_CASE_						:					Tuple			   =       [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_						:					Optional[int]			   =       time.time()
            SCREAMING_SNAKE_CASE_						:					List[str]			   =       fit(
                model,
                loss,
                optimizer,
                tf_train_dataset,
                args.epochs,
                args.per_device_train_batch_size,
                max_steps=args.max_steps,
            )
            SCREAMING_SNAKE_CASE_						:					str			   =       time.time() - start_train_time
            logger.info('''*** Train ***''')
            logger.info(F"train_runtime = {end_train_time}")

            SCREAMING_SNAKE_CASE_						:					List[str]			   =       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_						:					Any			   =       model.evaluate(tf_test_dataset, batch_size=args.per_device_eval_batch_size, return_dict=True)
            logger.info('''*** Evaluate ***''')

            SCREAMING_SNAKE_CASE_						:					Any			   =       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)












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


if is_torch_available():
 import torch

if is_vision_available():
 from PIL import Image

 from transformers import SwinaSRImageProcessor
 from transformers.image_transforms import get_image_size







class 		__UpperCAmelCase(		unittest.TestCase ):





     """simple docstring"""



     def __init__(      self :       Optional[int]						,				A :       Dict						,				A :       Union[str, Any]=7						,				A :       List[str]=3						,				A :       Optional[int]=18						,				A :       Optional[int]=30						,				A :       Optional[Any]=4_00						,				A :       Optional[Any]=True						,				A :       Tuple=1 / 2_55						,				A :       Tuple=True						,				A :       Optional[int]=8						,				)			->			Any:






         """simple docstring"""
         a__      =       parent
         a__      =       batch_size
         a__      =       num_channels
         a__      =       image_size
         a__      =       min_resolution
         a__      =       max_resolution
         a__      =       do_rescale
         a__      =       rescale_factor
         a__      =       do_pad
         a__      =       pad_size

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






         """simple docstring"""
         return {
             "do_rescale": self.do_rescale,
             "rescale_factor": self.rescale_factor,
             "do_pad": self.do_pad,
             "pad_size": self.pad_size,
         }






     def 				__SCREAMING_SNAKE_CASE		(      self :       int						,				A :       int=False						,				A :       Any=False						,				A :       List[Any]=False    )			->			Optional[Any]:






         """simple docstring"""
         assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

         if equal_resolution:
             a__      =       []
             for i in range(self.batch_size    ):
                 image_inputs.append(
                     np.random.randint(
                         2_55						,				size=(self.num_channels, self.max_resolution, self.max_resolution)						,				dtype=np.uinta    )    )
         else:
             a__      =       []
             for i in range(self.batch_size    ):
                 a__ ,		a__      =       np.random.choice(np.arange(self.min_resolution						,				self.max_resolution    )						,				2    )
                 image_inputs.append(np.random.randint(2_55						,				size=(self.num_channels, width, height)						,				dtype=np.uinta    )    )

         if not numpify and not torchify:
             # PIL expects the channel dimension as last dimension
             a__      =       [Image.fromarray(np.moveaxis(A						,				0						,				-1    )    ) for x in image_inputs]

         if torchify:
             a__      =       [torch.from_numpy(A    ) for x in image_inputs]

         return image_inputs







@require_torch
@require_vision
class 		__UpperCAmelCase(		lowercase__       ,	unittest.TestCase ):





     """simple docstring"""



     __UpperCAmelCase	: Dict						=    SwinaSRImageProcessor if is_vision_available() else None

     def 				__SCREAMING_SNAKE_CASE		(      self :       int    )			->			Dict:






         """simple docstring"""
         a__      =       SwinaSRImageProcessingTester(self    )

     @property
     def 				__SCREAMING_SNAKE_CASE		(      self :       List[Any]    )			->			List[Any]:






         """simple docstring"""
         return self.image_processor_tester.prepare_image_processor_dict()

     def 				__SCREAMING_SNAKE_CASE		(      self :       Tuple    )			->			int:






         """simple docstring"""
         a__      =       self.image_processing_class(**self.image_processor_dict    )
         self.assertTrue(hasattr(A						,				"""do_rescale"""    )    )
         self.assertTrue(hasattr(A						,				"""rescale_factor"""    )    )
         self.assertTrue(hasattr(A						,				"""do_pad"""    )    )
         self.assertTrue(hasattr(A						,				"""pad_size"""    )    )

     def 				__SCREAMING_SNAKE_CASE		(      self :       Dict    )			->			List[Any]:






         """simple docstring"""
         pass

     def 				__SCREAMING_SNAKE_CASE		(      self :       int						,				A :       Optional[Any]    )			->			List[str]:






         """simple docstring"""
         a__ ,		a__      =       get_image_size(A    )
         a__      =       self.image_processor_tester.pad_size

         a__      =       (old_height // size + 1) * size - old_height
         a__      =       (old_width // size + 1) * size - old_width
         return old_height + pad_height, old_width + pad_width

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






         """simple docstring"""
         # Initialize image_processor
         a__      =       self.image_processing_class(**self.image_processor_dict    )
         # create random PIL images
         a__      =       self.image_processor_tester.prepare_inputs(equal_resolution=A    )
         for image in image_inputs:
             self.assertIsInstance(A						,				Image.Image    )

         # Test not batched input
         a__      =       image_processor(image_inputs[0]						,				return_tensors="""pt"""    ).pixel_values
         a__ ,		a__      =       self.calculate_expected_size(np.array(image_inputs[0]    )    )
         self.assertEqual(
             encoded_images.shape						,				(
                 1,
                 self.image_processor_tester.num_channels,
                 expected_height,
                 expected_width,
             )						,				)

     def 				__SCREAMING_SNAKE_CASE		(      self :       Union[str, Any]    )			->			Optional[Any]:






         """simple docstring"""
         # Initialize image_processor
         a__      =       self.image_processing_class(**self.image_processor_dict    )
         # create random numpy tensors
         a__      =       self.image_processor_tester.prepare_inputs(equal_resolution=A						,				numpify=A    )
         for image in image_inputs:
             self.assertIsInstance(A						,				np.ndarray    )

         # Test not batched input
         a__      =       image_processor(image_inputs[0]						,				return_tensors="""pt"""    ).pixel_values
         a__ ,		a__      =       self.calculate_expected_size(image_inputs[0]    )
         self.assertEqual(
             encoded_images.shape						,				(
                 1,
                 self.image_processor_tester.num_channels,
                 expected_height,
                 expected_width,
             )						,				)






     def 				__SCREAMING_SNAKE_CASE		(      self :       Union[str, Any]    )			->			Tuple:






         """simple docstring"""
         # Initialize image_processor
         a__      =       self.image_processing_class(**self.image_processor_dict    )
         # create random PyTorch tensors
         a__      =       self.image_processor_tester.prepare_inputs(equal_resolution=A						,				torchify=A    )
         for image in image_inputs:
             self.assertIsInstance(A						,				torch.Tensor    )

         # Test not batched input
         a__      =       image_processor(image_inputs[0]						,				return_tensors="""pt"""    ).pixel_values
         a__ ,		a__      =       self.calculate_expected_size(image_inputs[0]    )
         self.assertEqual(
             encoded_images.shape						,				(
                 1,
                 self.image_processor_tester.num_channels,
                 expected_height,
                 expected_width,
             )						,				)






