





'''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():
    A_				     =							True
    A_				     =							"p3dn.24xlarge"
    import smdistributed.dataparallel.tensorflow as sdp
else:
    A_				     =							False




def       A							(      _UpperCAmelCase   :				Optional[int]						,_UpperCAmelCase   :				Optional[int]						,_UpperCAmelCase   :				Union[str, Any]						,_UpperCAmelCase   :				List[Any]						,_UpperCAmelCase   :				Union[str, Any]						,_UpperCAmelCase   :				List[str]						,_UpperCAmelCase   :				Dict=None					)  ->				Dict:


  '''simple docstring'''




  __lowerCAmelCase		:    Dict			     =  tqdm(_UpperCAmelCase					)
  for i, batch in enumerate(_UpperCAmelCase					):
    with tf.GradientTape() as tape:
      __lowerCAmelCase     ,					__lowerCAmelCase		:    Tuple			     =  batch
      __lowerCAmelCase		:    Optional[int]			     =  model(_UpperCAmelCase					)
      __lowerCAmelCase		:    Any			     =  loss(_UpperCAmelCase						,outputs.logits					)

    if SDP_ENABLED:
      __lowerCAmelCase		:    Any			     =  sdp.DistributedGradientTape(_UpperCAmelCase						,sparse_as_dense=_UpperCAmelCase					)

    __lowerCAmelCase		:    int			     =  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

  __lowerCAmelCase		:    str			     =  {'loss': loss_value.numpy()}
  return train_results




def       A							(      _UpperCAmelCase   :				int						,_UpperCAmelCase   :				int						,_UpperCAmelCase   :				str					)  ->				Dict:


  '''simple docstring'''




  __lowerCAmelCase     ,					__lowerCAmelCase		:    Union[str, Any]			     =  load_dataset('imdb'						,split=['train', 'test']					)

  # Preprocess train dataset
  __lowerCAmelCase		:    Any			     =  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']					)

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

  # Preprocess test dataset
  __lowerCAmelCase		:    List[str]			     =  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']					)

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

  if SDP_ENABLED:
    __lowerCAmelCase		:    int			     =  tf_train_dataset.shard(sdp.size()						,sdp.rank()					)
    __lowerCAmelCase		:    Optional[Any]			     =  tf_test_dataset.shard(sdp.size()						,sdp.rank()					)
  __lowerCAmelCase		:    Optional[Any]			     =  tf_train_dataset.batch(_UpperCAmelCase						,drop_remainder=_UpperCAmelCase					)
  __lowerCAmelCase		:    List[str]			     =  tf_test_dataset.batch(_UpperCAmelCase						,drop_remainder=_UpperCAmelCase					)

  return tf_train_dataset, tf_test_dataset


if __name__ == "__main__":
    A_				     =							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"])

    A_     ,  A_				     =							parser.parse_known_args()

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

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

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

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

        A_				     =							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 typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


A_				     =							{
    "configuration_sam": [
        "SAM_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "SamConfig",
        "SamMaskDecoderConfig",
        "SamPromptEncoderConfig",
        "SamVisionConfig",
    ],
    "processing_sam": ["SamProcessor"],
}


try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    A_				     =							[
        "SAM_PRETRAINED_MODEL_ARCHIVE_LIST",
        "SamModel",
        "SamPreTrainedModel",
    ]
try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    A_				     =							[
        "TF_SAM_PRETRAINED_MODEL_ARCHIVE_LIST",
        "TFSamModel",
        "TFSamPreTrainedModel",
    ]
try:
    if not is_vision_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    A_				     =							["SamImageProcessor"]


if TYPE_CHECKING:
    from .configuration_sam import (
        SAM_PRETRAINED_CONFIG_ARCHIVE_MAP,
        SamConfig,
        SamMaskDecoderConfig,
        SamPromptEncoderConfig,
        SamVisionConfig,
    )
    from .processing_sam import SamProcessor

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_sam import SAM_PRETRAINED_MODEL_ARCHIVE_LIST, SamModel, SamPreTrainedModel

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_sam import TF_SAM_PRETRAINED_MODEL_ARCHIVE_LIST, TFSamModel, TFSamPreTrainedModel

    try:
        if not is_vision_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .image_processing_sam import SamImageProcessor

else:
    import sys

    A_				     =							_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





