
"""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():
      UpperCamelCase_						   =   True
      UpperCamelCase_						   =   'p3dn.24xlarge'
      import smdistributed.dataparallel.tensorflow as sdp
else:
      UpperCamelCase_						   =   False
def    UpperCamelCase						(					UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase=None     )						->Optional[Any]:

   """simple docstring"""



   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    UpperCamelCase						(					UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase     )						->Dict:

   """simple docstring"""



   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__":
      UpperCamelCase_						   =   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'])

      UpperCamelCase_   ,     UpperCamelCase_						   =   parser.parse_known_args()

      # Set up logging
      UpperCamelCase_						   =   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()

            UpperCamelCase_						   =   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
      UpperCamelCase_						   =   TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
      UpperCamelCase_						   =   AutoTokenizer.from_pretrained(args.model_name_or_path)

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

            UpperCamelCase_						   =   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):
            UpperCamelCase_						   =   model.evaluate(tf_test_dataset, batch_size=args.per_device_eval_batch_size, return_dict=True)
            logger.info('*** Evaluate ***')

            UpperCamelCase_						   =   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"""






from collections import namedtuple

UpperCamelCase_						   =   namedtuple('from_to', 'from_ to')

UpperCamelCase_						   =   {
    'atm': from_to(1, 1),
    'pascal': from_to(0.0_00_00_98, 101325),
    'bar': from_to(0.98_69_23, 1.0_13_25),
    'kilopascal': from_to(0.00_98_69_23, 1_01.3_25),
    'megapascal': from_to(9.8_69_23, 0.10_13_25),
    'psi': from_to(0.06_80_46, 14.69_59),
    'inHg': from_to(0.0_33_42_11, 29.92_13),
    'torr': from_to(0.00_13_15_79, 760),
}
def    UpperCamelCase						(					UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase     )						->float:

   """simple docstring"""



   if from_type not in PRESSURE_CONVERSION:
      raise ValueError(
          F'''Invalid \'from_type\' value: {from_type!r}  Supported values are:\n'''
          + ", ".join(UpperCAmelCase     )     )
   if to_type not in PRESSURE_CONVERSION:
      raise ValueError(
          F'''Invalid \'to_type\' value: {to_type!r}.  Supported values are:\n'''
          + ", ".join(UpperCAmelCase     )     )
   return (
       value * PRESSURE_CONVERSION[from_type].from_ * PRESSURE_CONVERSION[to_type].to
   )


if __name__ == "__main__":
      import doctest

      doctest.testmod()