



'''simple docstring'''







import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch
      from torch import nn

      from transformers import DinatBackbone, DinatForImageClassification, DinatModel
      from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor

class    lowerCamelCase_				:






   '''simple docstring'''

   def __init__(	self  ,    __UpperCamelCase  ,    __UpperCamelCase=13  ,    __UpperCamelCase=64  ,    __UpperCamelCase=4  ,    __UpperCamelCase=3  ,    __UpperCamelCase=16  ,    __UpperCamelCase=[1, 2, 1]  ,    __UpperCamelCase=[2, 4, 8]  ,    __UpperCamelCase=3  ,    __UpperCamelCase=[[3], [1, 2], [1]]  ,    __UpperCamelCase=2.0  ,    __UpperCamelCase=True  ,    __UpperCamelCase=0.0  ,    __UpperCamelCase=0.0  ,    __UpperCamelCase=0.1  ,    __UpperCamelCase="gelu"  ,    __UpperCamelCase=True  ,    __UpperCamelCase=0.02  ,    __UpperCamelCase=1e-5  ,    __UpperCamelCase=True  ,    __UpperCamelCase=None  ,    __UpperCamelCase=True  ,    __UpperCamelCase=10  ,    __UpperCamelCase=["stage1", "stage2"]  ,    __UpperCamelCase=[1, 2]  ,    ):
       UpperCamelCase__	     =      parent
       UpperCamelCase__	     =      batch_size
       UpperCamelCase__	     =      image_size
       UpperCamelCase__	     =      patch_size
       UpperCamelCase__	     =      num_channels
       UpperCamelCase__	     =      embed_dim
       UpperCamelCase__	     =      depths
       UpperCamelCase__	     =      num_heads
       UpperCamelCase__	     =      kernel_size
       UpperCamelCase__	     =      dilations
       UpperCamelCase__	     =      mlp_ratio
       UpperCamelCase__	     =      qkv_bias
       UpperCamelCase__	     =      hidden_dropout_prob
       UpperCamelCase__	     =      attention_probs_dropout_prob
       UpperCamelCase__	     =      drop_path_rate
       UpperCamelCase__	     =      hidden_act
       UpperCamelCase__	     =      patch_norm
       UpperCamelCase__	     =      layer_norm_eps
       UpperCamelCase__	     =      initializer_range
       UpperCamelCase__	     =      is_training
       UpperCamelCase__	     =      scope
       UpperCamelCase__	     =      use_labels
       UpperCamelCase__	     =      num_labels
       UpperCamelCase__	     =      out_features
       UpperCamelCase__	     =      out_indices


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

       UpperCamelCase__	     =      None
       if self.use_labels:
           UpperCamelCase__	     =      ids_tensor([self.batch_size]  ,    self.num_labels			)

       UpperCamelCase__	     =      self.get_config()

       return config, pixel_values, labels


   def 						_UpperCamelCase							(	self			):
       return DinatConfig(
           num_labels=self.num_labels  ,    image_size=self.image_size  ,    patch_size=self.patch_size  ,    num_channels=self.num_channels  ,    embed_dim=self.embed_dim  ,    depths=self.depths  ,    num_heads=self.num_heads  ,    kernel_size=self.kernel_size  ,    dilations=self.dilations  ,    mlp_ratio=self.mlp_ratio  ,    qkv_bias=self.qkv_bias  ,    hidden_dropout_prob=self.hidden_dropout_prob  ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,    drop_path_rate=self.drop_path_rate  ,    hidden_act=self.hidden_act  ,    patch_norm=self.patch_norm  ,    layer_norm_eps=self.layer_norm_eps  ,    initializer_range=self.initializer_range  ,    out_features=self.out_features  ,    out_indices=self.out_indices  ,    )


   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			):
       UpperCamelCase__	     =      DinatModel(config=__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()
       UpperCamelCase__	     =      model(__UpperCamelCase			)

       UpperCamelCase__	     =      UpperCamelCase__	     =      (config.image_size // config.patch_size) // (2 ** (len(config.depths			) - 1))
       UpperCamelCase__	     =      int(config.embed_dim * 2 ** (len(config.depths			) - 1)			)

       self.parent.assertEqual(
           result.last_hidden_state.shape  ,    (self.batch_size, expected_height, expected_width, expected_dim)			)


   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			):
       UpperCamelCase__	     =      DinatForImageClassification(__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()
       UpperCamelCase__	     =      model(__UpperCamelCase  ,    labels=__UpperCamelCase			)
       self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.num_labels)			)

       # test greyscale images
       UpperCamelCase__	     =      1
       UpperCamelCase__	     =      DinatForImageClassification(__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()

       UpperCamelCase__	     =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
       UpperCamelCase__	     =      model(__UpperCamelCase			)
       self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.num_labels)			)


   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			):
       UpperCamelCase__	     =      DinatBackbone(config=__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()
       UpperCamelCase__	     =      model(__UpperCamelCase			)

       # verify hidden states
       self.parent.assertEqual(len(result.feature_maps			)  ,    len(config.out_features			)			)
       self.parent.assertListEqual(list(result.feature_maps[0].shape			)  ,    [self.batch_size, model.channels[0], 16, 16]			)

       # verify channels
       self.parent.assertEqual(len(model.channels			)  ,    len(config.out_features			)			)

       # verify backbone works with out_features=None
       UpperCamelCase__	     =      None
       UpperCamelCase__	     =      DinatBackbone(config=__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()
       UpperCamelCase__	     =      model(__UpperCamelCase			)

       # verify feature maps
       self.parent.assertEqual(len(result.feature_maps			)  ,    1			)
       self.parent.assertListEqual(list(result.feature_maps[0].shape			)  ,    [self.batch_size, model.channels[-1], 4, 4]			)

       # verify channels
       self.parent.assertEqual(len(model.channels			)  ,    1			)







   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      self.prepare_config_and_inputs()
       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      config_and_inputs
       UpperCamelCase__	     =      {"pixel_values": pixel_values}
       return config, inputs_dict

@require_natten
@require_torch
class    lowerCamelCase_				(						lowerCAmelCase_       ,     lowerCAmelCase_       ,     unittest.TestCase	):






   '''simple docstring'''

   A			:  int					=						(
       (
           DinatModel,
           DinatForImageClassification,
           DinatBackbone,
       )
       if is_torch_available()
       else ()
   )
   A			:  Tuple					=						(
       {"feature-extraction": DinatModel, "image-classification": DinatForImageClassification}
       if is_torch_available()
       else {}
   )
   A			:  Optional[Any]					=						False

   A			:  List[str]					=						False
   A			:  Any					=						False
   A			:  Optional[Any]					=						False
   A			:  List[Any]					=						False


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      DinatModelTester(self			)
       UpperCamelCase__	     =      ConfigTester(self  ,    config_class=__UpperCamelCase  ,    embed_dim=37			)


   def 						_UpperCamelCase							(	self			):
       self.create_and_test_config_common_properties()
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()


   def 						_UpperCamelCase							(	self			):
       return


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__UpperCamelCase			)


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase			)


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_backbone(*__UpperCamelCase			)


   @unittest.skip(reason="Dinat does not use inputs_embeds"			)
   def 						_UpperCamelCase							(	self			):
       pass


   @unittest.skip(reason="Dinat does not use feedforward chunking"			)
   def 						_UpperCamelCase							(	self			):
       pass


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           UpperCamelCase__	     =      model_class(__UpperCamelCase			)
           self.assertIsInstance(model.get_input_embeddings()  ,    (nn.Module)			)
           UpperCamelCase__	     =      model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(__UpperCamelCase  ,    nn.Linear			)			)


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           UpperCamelCase__	     =      model_class(__UpperCamelCase			)
           UpperCamelCase__	     =      inspect.signature(model.forward			)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           UpperCamelCase__	     =      [*signature.parameters.keys()]

           UpperCamelCase__	     =      ["pixel_values"]
           self.assertListEqual(arg_names[:1]  ,    __UpperCamelCase			)


   def 						_UpperCamelCase							(	self			):
       self.skipTest("Dinat's attention operation is handled entirely by NATTEN."			)


   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			):
       UpperCamelCase__	     =      model_class(__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()

       with torch.no_grad():
           UpperCamelCase__	     =      model(**self._prepare_for_class(__UpperCamelCase  ,    __UpperCamelCase			)			)

       UpperCamelCase__	     =      outputs.hidden_states

       UpperCamelCase__	     =      getattr(
           self.model_tester  ,    "expected_num_hidden_layers"  ,    len(self.model_tester.depths			) + 1			)
       self.assertEqual(len(__UpperCamelCase			)  ,    __UpperCamelCase			)

       # Dinat has a different seq_length
       UpperCamelCase__	     =      (
           config.patch_size
           if isinstance(config.patch_size  ,    collections.abc.Iterable			)
           else (config.patch_size, config.patch_size)
       )

       UpperCamelCase__	     =      image_size[0] // patch_size[0]
       UpperCamelCase__	     =      image_size[1] // patch_size[1]

       self.assertListEqual(
           list(hidden_states[0].shape[-3:]			)  ,    [height, width, self.model_tester.embed_dim]  ,    )

       if model_class.__name__ != "DinatBackbone":
           UpperCamelCase__	     =      outputs.reshaped_hidden_states
           self.assertEqual(len(__UpperCamelCase			)  ,    __UpperCamelCase			)

           UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      reshaped_hidden_states[0].shape
           UpperCamelCase__	     =      (
               reshaped_hidden_states[0].view(__UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			).permute(0  ,    2  ,    3  ,    1			)
           )
           self.assertListEqual(
               list(reshaped_hidden_states.shape[-3:]			)  ,    [height, width, self.model_tester.embed_dim]  ,    )


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.model_tester.prepare_config_and_inputs_for_common()

       UpperCamelCase__	     =      (
           self.model_tester.image_size
           if isinstance(self.model_tester.image_size  ,    collections.abc.Iterable			)
           else (self.model_tester.image_size, self.model_tester.image_size)
       )

       for model_class in self.all_model_classes:
           UpperCamelCase__	     =      True
           self.check_hidden_states_output(__UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           UpperCamelCase__	     =      True

           self.check_hidden_states_output(__UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			)


   @slow
   def 						_UpperCamelCase							(	self			):
       for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCamelCase__	     =      DinatModel.from_pretrained(__UpperCamelCase			)
           self.assertIsNotNone(__UpperCamelCase			)







   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.model_tester.prepare_config_and_inputs_for_common()

       UpperCamelCase__	     =      _config_zero_init(__UpperCamelCase			)
       for model_class in self.all_model_classes:
           UpperCamelCase__	     =      model_class(config=__UpperCamelCase			)
           for name, param in model.named_parameters():
               if "embeddings" not in name and param.requires_grad:
                   self.assertIn(
                       ((param.data.mean() * 1e9).round() / 1e9).item()  ,    [0.0, 1.0]  ,    msg=f'Parameter {name} of model {model_class} seems not properly initialized'  ,    )

@require_natten
@require_vision
@require_torch
class    lowerCamelCase_				(						unittest.TestCase	):






   '''simple docstring'''

   @cached_property
   def 						_UpperCamelCase							(	self			):
       return AutoImageProcessor.from_pretrained("shi-labs/dinat-mini-in1k-224"			) if is_vision_available() else None







   @slow
   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      DinatForImageClassification.from_pretrained("shi-labs/dinat-mini-in1k-224"			).to(__UpperCamelCase			)
       UpperCamelCase__	     =      self.default_image_processor

       UpperCamelCase__	     =      Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
       UpperCamelCase__	     =      image_processor(images=__UpperCamelCase  ,    return_tensors="pt"			).to(__UpperCamelCase			)

       # forward pass
       with torch.no_grad():
           UpperCamelCase__	     =      model(**__UpperCamelCase			)

       # verify the logits
       UpperCamelCase__	     =      torch.Size((1, 10_00)			)
       self.assertEqual(outputs.logits.shape  ,    __UpperCamelCase			)
       UpperCamelCase__	     =      torch.tensor([-0.1_545, -0.7_667, 0.4_642]			).to(__UpperCamelCase			)
       self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,    __UpperCamelCase  ,    atol=1e-4			)			)

@require_torch
@require_natten
class    lowerCamelCase_				(						unittest.TestCase       ,     lowerCAmelCase_	):






   '''simple docstring'''

   A			:  Optional[Any]					=						(DinatBackbone,) if is_torch_available() else ()
   A			:  Tuple					=						DinatConfig


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      DinatModelTester(self			)









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





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_=None						):
    UpperCamelCase__	     =      tqdm(lowerCAmelCase_						)
    for i, batch in enumerate(lowerCAmelCase_						):
        with tf.GradientTape() as tape:
            UpperCamelCase__     ,	UpperCamelCase__	     =      batch
            UpperCamelCase__	     =      model(lowerCAmelCase_						)
            UpperCamelCase__	     =      loss(lowerCAmelCase_       ,	outputs.logits						)

        if SDP_ENABLED:
            UpperCamelCase__	     =      sdp.DistributedGradientTape(lowerCAmelCase_       ,	sparse_as_dense=lowerCAmelCase_						)

        UpperCamelCase__	     =      tape.gradient(lowerCAmelCase_       ,	model.trainable_variables						)
        opt.apply_gradients(zip(lowerCAmelCase_       ,	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

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





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_						):
    # Load dataset
    UpperCamelCase__     ,	UpperCamelCase__	     =      load_dataset("imdb"       ,	split=["train", "test"]						)

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

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

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

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

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

    return tf_train_dataset, tf_test_dataset


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

      lowercase__       ,	lowercase__  :	Union[str, Any]            =      parser.parse_known_args()

      # Set up logging
      lowercase__  :	int            =      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()

            lowercase__  :	Union[str, 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
      lowercase__  :	int            =      TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
      lowercase__  :	Any            =      AutoTokenizer.from_pretrained(args.model_name_or_path)

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

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

            lowercase__  :	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)




