'''simple docstring'''





import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

UpperCamelCase    :			Any						     =  logging.get_logger(__name__)






def SCREAMING_SNAKE_CASE__   (							snake_case    :							List[Any]							,     snake_case    :							bool = True   )					->							Tuple:







     """simple docstring"""
     a :				Optional[Any]				=      {}
     for name in model.attn_processors.keys():
          a :				Optional[Any]				=      name.endswith('attn2.processor'   ) and not (
              name.startswith('transformer_in'   ) or 'temp_attentions' in name.split('.'   )
          )
          a :				int				=      model.config.cross_attention_dim if has_cross_attention else None
          if name.startswith('mid_block'   ):
               a :				List[Any]				=      model.config.block_out_channels[-1]
          elif name.startswith('up_blocks'   ):
               a :				int				=      int(name[len('up_blocks.'   )]   )
               a :				Tuple				=      list(reversed(model.config.block_out_channels   )   )[block_id]
          elif name.startswith('down_blocks'   ):
               a :				List[str]				=      int(name[len('down_blocks.'   )]   )
               a :				Dict				=      model.config.block_out_channels[block_id]
          elif name.startswith('transformer_in'   ):
               # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
               a :				Union[str, Any]				=      8 * model.config.attention_head_dim

          a :				Optional[int]				=      LoRAAttnProcessor(hidden_size=lowerCAmelCase_							,     cross_attention_dim=lowerCAmelCase_   )
          a :				str				=      lora_attn_procs[name].to(model.device   )

          if mock_weights:
               # add 1 to weights to mock trained weights
               with torch.no_grad():
                    lora_attn_procs[name].to_q_lora.up.weight += 1
                    lora_attn_procs[name].to_k_lora.up.weight += 1
                    lora_attn_procs[name].to_v_lora.up.weight += 1
                    lora_attn_procs[name].to_out_lora.up.weight += 1

     return lora_attn_procs





@skip_mps
class    UpperCamelCase (    snake_case__  ,			snake_case__  ,			unittest.TestCase    ):







    """simple docstring"""





    A			:  Union[str, Any]            =       UNetaDConditionModel
    A			:  List[Any]            =       "sample"





    @property
    def 					SCREAMING_SNAKE_CASE_			( self			:		str):


         """simple docstring"""







         a :				Optional[Any]				=      4
         a :				str				=      4
         a :				List[str]				=      4
         a :				str				=      (3_2, 3_2)

         a :				List[Any]				=      floats_tensor((batch_size, num_channels, num_frames) + sizes).to(__UpperCAmelCase)
         a :				str				=      torch.tensor([1_0]).to(__UpperCAmelCase)
         a :				str				=      floats_tensor((batch_size, 4, 3_2)).to(__UpperCAmelCase)

         return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}





    @property
    def 					SCREAMING_SNAKE_CASE_			( self			:		str):


         """simple docstring"""







         return (4, 4, 3_2, 3_2)





    @property
    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[int]):


         """simple docstring"""







         return (4, 4, 3_2, 3_2)





    def 					SCREAMING_SNAKE_CASE_			( self			:		Dict):


         """simple docstring"""







         a :				Optional[int]				=      {
             'block_out_channels': (3_2, 6_4),
             'down_block_types': (
                 'CrossAttnDownBlock3D',
                 'DownBlock3D',
             ),
             'up_block_types': ('UpBlock3D', 'CrossAttnUpBlock3D'),
             'cross_attention_dim': 3_2,
             'attention_head_dim': 8,
             'out_channels': 4,
             'in_channels': 4,
             'layers_per_block': 1,
             'sample_size': 3_2,
         }
         a :				Optional[Any]				=      self.dummy_input
         return init_dict, inputs_dict





    @unittest.skipIf(
        torch_device != 'cuda' or not is_xformers_available()				,				reason='XFormers attention is only available with CUDA and `xformers` installed'				,				)
    def 					SCREAMING_SNAKE_CASE_			( self			:		Tuple):


         """simple docstring"""







         a				,       a :				Any				=      self.prepare_init_args_and_inputs_for_common()
         a :				Optional[Any]				=      self.model_class(**__UpperCAmelCase)

         model.enable_xformers_memory_efficient_attention()

         assert (
             model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
             == "XFormersAttnProcessor"
         ), "xformers is not enabled"





    def 					SCREAMING_SNAKE_CASE_			( self			:		Union[str, Any]):


         """simple docstring"""







         a				,       a :				Any				=      self.prepare_init_args_and_inputs_for_common()

         a :				List[Any]				=      3_2

         a :				List[str]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)
         model.eval()

         with torch.no_grad():
              a :				Dict				=      model(**__UpperCAmelCase)

              if isinstance(__UpperCAmelCase				,				__UpperCAmelCase):
                   a :				Optional[Any]				=      output.sample

         self.assertIsNotNone(__UpperCAmelCase)
         a :				Union[str, Any]				=      inputs_dict['sample'].shape
         self.assertEqual(output.shape				,				__UpperCAmelCase				,				'Input and output shapes do not match')





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[Any]):


         """simple docstring"""







         a				,       a :				int				=      self.prepare_init_args_and_inputs_for_common()
         a :				List[Any]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)
         model.eval()

         with torch.no_grad():
              # Warmup pass when using mps (see #372)
              if torch_device == "mps" and isinstance(__UpperCAmelCase				,				__UpperCAmelCase):
                   model(**self.dummy_input)

              a :				List[Any]				=      model(**__UpperCAmelCase)
              if isinstance(__UpperCAmelCase				,				__UpperCAmelCase):
                   a :				Union[str, Any]				=      first.sample

              a :				Optional[Any]				=      model(**__UpperCAmelCase)
              if isinstance(__UpperCAmelCase				,				__UpperCAmelCase):
                   a :				Optional[Any]				=      second.sample

         a :				List[Any]				=      first.cpu().numpy()
         a :				Optional[Any]				=      second.cpu().numpy()
         a :				Any				=      out_a[~np.isnan(__UpperCAmelCase)]
         a :				Optional[int]				=      out_a[~np.isnan(__UpperCAmelCase)]
         a :				int				=      np.amax(np.abs(out_a - out_a))
         self.assertLessEqual(__UpperCAmelCase				,				1e-5)





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[int]):


         """simple docstring"""







         a				,       a :				str				=      self.prepare_init_args_and_inputs_for_common()

         a :				Any				=      8

         a :				List[str]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)
         model.eval()

         model.set_attention_slice('auto')
         with torch.no_grad():
              a :				Optional[int]				=      model(**__UpperCAmelCase)
         assert output is not None

         model.set_attention_slice('max')
         with torch.no_grad():
              a :				str				=      model(**__UpperCAmelCase)
         assert output is not None

         model.set_attention_slice(2)
         with torch.no_grad():
              a :				List[Any]				=      model(**__UpperCAmelCase)
         assert output is not None





    def 					SCREAMING_SNAKE_CASE_			( self			:		Any):


         """simple docstring"""







         a				,       a :				Union[str, Any]				=      self.prepare_init_args_and_inputs_for_common()

         a :				Any				=      8

         a :				Any				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)

         with torch.no_grad():
              a :				Optional[Any]				=      model(**__UpperCAmelCase).sample

         a :				List[Any]				=      create_lora_layers(__UpperCAmelCase)

         # make sure we can set a list of attention processors
         model.set_attn_processor(__UpperCAmelCase)
         model.to(__UpperCAmelCase)

         # test that attn processors can be set to itself
         model.set_attn_processor(model.attn_processors)

         with torch.no_grad():
              a :				Optional[int]				=      model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.0}).sample
              a :				Dict				=      model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.5}).sample
              a :				Union[str, Any]				=      model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.5}).sample

         assert (samplea - samplea).abs().max() < 3e-3
         assert (samplea - samplea).abs().max() < 3e-3

         # sample 2 and sample 3 should be different
         assert (samplea - samplea).abs().max() > 3e-3





    def 					SCREAMING_SNAKE_CASE_			( self			:		Tuple):


         """simple docstring"""







         a				,       a :				List[str]				=      self.prepare_init_args_and_inputs_for_common()

         a :				List[Any]				=      8

         torch.manual_seed(0)
         a :				Optional[int]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)

         with torch.no_grad():
              a :				Optional[int]				=      model(**__UpperCAmelCase).sample

         a :				List[str]				=      create_lora_layers(__UpperCAmelCase)
         model.set_attn_processor(__UpperCAmelCase)

         with torch.no_grad():
              a :				Optional[Any]				=      model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.5}).sample

         with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_attn_procs(__UpperCAmelCase)
              self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase				,				'pytorch_lora_weights.bin')))
              torch.manual_seed(0)
              a :				Tuple				=      self.model_class(**__UpperCAmelCase)
              new_model.to(__UpperCAmelCase)
              new_model.load_attn_procs(__UpperCAmelCase)

         with torch.no_grad():
              a :				Any				=      new_model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.5}).sample

         assert (sample - new_sample).abs().max() < 1e-3

         # LoRA and no LoRA should NOT be the same
         assert (sample - old_sample).abs().max() > 1e-4





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[Any]):


         """simple docstring"""







         a				,       a :				List[str]				=      self.prepare_init_args_and_inputs_for_common()

         a :				Optional[int]				=      8

         torch.manual_seed(0)
         a :				List[str]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)

         with torch.no_grad():
              a :				Optional[int]				=      model(**__UpperCAmelCase).sample

         a :				Optional[int]				=      create_lora_layers(__UpperCAmelCase)
         model.set_attn_processor(__UpperCAmelCase)

         with torch.no_grad():
              a :				List[str]				=      model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.5}).sample

         with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_attn_procs(__UpperCAmelCase				,				safe_serialization=__UpperCAmelCase)
              self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase				,				'pytorch_lora_weights.safetensors')))
              torch.manual_seed(0)
              a :				Tuple				=      self.model_class(**__UpperCAmelCase)
              new_model.to(__UpperCAmelCase)
              new_model.load_attn_procs(__UpperCAmelCase)

         with torch.no_grad():
              a :				Any				=      new_model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.5}).sample

         assert (sample - new_sample).abs().max() < 3e-3

         # LoRA and no LoRA should NOT be the same
         assert (sample - old_sample).abs().max() > 1e-4





    def 					SCREAMING_SNAKE_CASE_			( self			:		List[Any]):


         """simple docstring"""







         a				,       a :				int				=      self.prepare_init_args_and_inputs_for_common()

         a :				int				=      8

         torch.manual_seed(0)
         a :				Any				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)

         a :				int				=      create_lora_layers(__UpperCAmelCase				,				mock_weights=__UpperCAmelCase)
         model.set_attn_processor(__UpperCAmelCase)
         # Saving as torch, properly reloads with directly filename
         with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_attn_procs(__UpperCAmelCase)
              self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase				,				'pytorch_lora_weights.bin')))
              torch.manual_seed(0)
              a :				List[Any]				=      self.model_class(**__UpperCAmelCase)
              new_model.to(__UpperCAmelCase)
              new_model.load_attn_procs(__UpperCAmelCase				,				weight_name='pytorch_lora_weights.bin')





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[Any]):


         """simple docstring"""







         a				,       a :				Dict				=      self.prepare_init_args_and_inputs_for_common()

         a :				Optional[Any]				=      8

         torch.manual_seed(0)
         a :				List[Any]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)

         a :				Optional[int]				=      create_lora_layers(__UpperCAmelCase				,				mock_weights=__UpperCAmelCase)
         model.set_attn_processor(__UpperCAmelCase)
         # Saving as torch, properly reloads with directly filename
         with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_attn_procs(__UpperCAmelCase)
              self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase				,				'pytorch_lora_weights.bin')))
              torch.manual_seed(0)
              a :				Tuple				=      self.model_class(**__UpperCAmelCase)
              new_model.to(__UpperCAmelCase)
              with self.assertRaises(__UpperCAmelCase) as e:
                   new_model.load_attn_procs(__UpperCAmelCase				,				use_safetensors=__UpperCAmelCase)
              self.assertIn('Error no file named pytorch_lora_weights.safetensors'				,				str(e.exception))





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[Any]):


         """simple docstring"""







         a				,       a :				int				=      self.prepare_init_args_and_inputs_for_common()

         a :				Union[str, Any]				=      8

         torch.manual_seed(0)
         a :				List[str]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)

         with torch.no_grad():
              a :				List[Any]				=      model(**__UpperCAmelCase).sample

         a :				List[str]				=      create_lora_layers(__UpperCAmelCase)
         model.set_attn_processor(__UpperCAmelCase)

         with torch.no_grad():
              a :				Optional[int]				=      model(**__UpperCAmelCase				,				cross_attention_kwargs={'scale': 0.0}).sample

         model.set_attn_processor(AttnProcessor())

         with torch.no_grad():
              a :				Tuple				=      model(**__UpperCAmelCase).sample

         assert (sample - new_sample).abs().max() < 1e-4
         assert (sample - old_sample).abs().max() < 3e-3





    @unittest.skipIf(
        torch_device != 'cuda' or not is_xformers_available()				,				reason='XFormers attention is only available with CUDA and `xformers` installed'				,				)
    def 					SCREAMING_SNAKE_CASE_			( self			:		Union[str, Any]):


         """simple docstring"""







         a				,       a :				Any				=      self.prepare_init_args_and_inputs_for_common()

         a :				Optional[Any]				=      4

         torch.manual_seed(0)
         a :				Dict				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)
         a :				Tuple				=      create_lora_layers(__UpperCAmelCase)
         model.set_attn_processor(__UpperCAmelCase)

         # default
         with torch.no_grad():
              a :				List[Any]				=      model(**__UpperCAmelCase).sample

              model.enable_xformers_memory_efficient_attention()
              a :				int				=      model(**__UpperCAmelCase).sample

              model.disable_xformers_memory_efficient_attention()
              a :				Union[str, Any]				=      model(**__UpperCAmelCase).sample

         assert (sample - on_sample).abs().max() < 1e-4
         assert (sample - off_sample).abs().max() < 1e-4







    def 					SCREAMING_SNAKE_CASE_			( self			:		List[Any]):


         """simple docstring"""







         a				,       a :				Optional[Any]				=      self.prepare_init_args_and_inputs_for_common()
         a :				Dict				=      3_2

         a :				Union[str, Any]				=      self.model_class(**__UpperCAmelCase)
         model.to(__UpperCAmelCase)
         model.eval()

         with torch.no_grad():
              a :				Optional[Any]				=      model(**__UpperCAmelCase)[0]

         model.enable_forward_chunking()
         with torch.no_grad():
              a :				Union[str, Any]				=      model(**__UpperCAmelCase)[0]

         self.assertEqual(output.shape				,				output_a.shape				,				'Shape doesn\'t match')
         assert np.abs(output.cpu() - output_a.cpu()).max() < 1e-2


# (todo: sayakpaul) implement SLOW tests.






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     (    lowerCAmelCase_   :					Optional[Any]			,lowerCAmelCase_   :					Union[str, Any]			,lowerCAmelCase_   :					int			,lowerCAmelCase_   :					int			,lowerCAmelCase_   :					Optional[int]			,lowerCAmelCase_   :					Any			,lowerCAmelCase_   :					int=None				)      -> Optional[int]:






						'''simple docstring'''



						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        __a     (    lowerCAmelCase_   :					Optional[int]			,lowerCAmelCase_   :					Dict			,lowerCAmelCase_   :					str				)      -> str:






						'''simple docstring'''



						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__":
	__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)




