






"""simple docstring"""


import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   import jax
   import jax.numpy as jnp
   import numpy as np

   from transformers.modeling_flax_pytorch_utils import (
       convert_pytorch_state_dict_to_flax,
       load_flax_weights_in_pytorch_model,
   )
   from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
   import torch






@require_flax
class 						snake_case  :





    def __init__(							self	:					str				,		A	:					List[str]				,		A	:					Any=1_4				,		A	:					str=7				,		A	:					Dict=True				,		A	:					Optional[int]=True				,		A	:					str=True				,		A	:					List[Any]=9_9				,		A	:					Union[str, Any]=3_2				,		A	:					Any=5				,		A	:					Tuple=4				,		A	:					Any=3_7				,		A	:					Dict="gelu"				,		A	:					Tuple=0.1				,		A	:					List[str]=0.1				,		A	:					Any=5_1_2				,		A	:					Tuple=0.02				,		A	:					str=None				,		):


      '''simple docstring'''




      a     :			str				=   parent
      a     :			Union[str, Any]				=   batch_size
      a     :			str				=   seq_length
      a     :			List[Any]				=   is_training
      a     :			Any				=   use_input_mask
      a     :			Dict				=   use_labels
      a     :			int				=   vocab_size
      a     :			str				=   d_model
      a     :			Union[str, Any]				=   num_hidden_layers
      a     :			Optional[int]				=   num_attention_heads
      a     :			List[str]				=   ffn_dim
      a     :			List[Any]				=   activation_function
      a     :			str				=   activation_dropout
      a     :			Optional[Any]				=   attention_dropout
      a     :			Tuple				=   max_position_embeddings
      a     :			str				=   initializer_range
      a     :			List[str]				=   None
      a     :			Optional[Any]				=   0
      a     :			Tuple				=   2
      a     :			List[str]				=   1





    def 				lowerCamelCase__      (							self	:					Tuple			):


      '''simple docstring'''




      a     :			Optional[Any]				=   np.clip(ids_tensor([self.batch_size, self.seq_length]				,		self.vocab_size			)				,		3				,		self.vocab_size			)

      a     :			Optional[int]				=   None
      if self.use_input_mask:
        a     :			List[str]				=   random_attention_mask([self.batch_size, self.seq_length]			)

      a     :			List[str]				=   XGLMConfig(
          vocab_size=self.vocab_size				,		d_model=self.hidden_size				,		num_layers=self.num_hidden_layers				,		attention_heads=self.num_attention_heads				,		ffn_dim=self.ffn_dim				,		activation_function=self.activation_function				,		activation_dropout=self.activation_dropout				,		attention_dropout=self.attention_dropout				,		max_position_embeddings=self.max_position_embeddings				,		initializer_range=self.initializer_range				,		use_cache=__UpperCAmelCase				,		bos_token_id=self.bos_token_id				,		eos_token_id=self.eos_token_id				,		pad_token_id=self.pad_token_id				,		)

      return (config, input_ids, input_mask)





    def 				lowerCamelCase__      (							self	:					List[str]			):


      '''simple docstring'''




      a     :			Union[str, Any]				=   self.prepare_config_and_inputs()
      a,     a,     a     :			Tuple				=   config_and_inputs
      a     :			Optional[Any]				=   {'input_ids': input_ids, 'attention_mask': attention_mask}
      return config, inputs_dict





    def 				lowerCamelCase__      (							self	:					Any			):


      '''simple docstring'''




      a,     a,     a     :			Optional[int]				=   self.prepare_config_and_inputs()

      a     :			Optional[int]				=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
      a     :			Tuple				=   ids_tensor([self.batch_size, self.seq_length]				,		vocab_size=2			)

      return (
          config,
          input_ids,
          attention_mask,
          encoder_hidden_states,
          encoder_attention_mask,
      )





    def 				lowerCamelCase__      (							self	:					Tuple				,		A	:					Optional[Any]				,		A	:					Optional[int]				,		A	:					List[str]				,		A	:					Tuple			):


      '''simple docstring'''




      a     :			Any				=   2_0
      a     :			Union[str, Any]				=   model_class_name(__UpperCAmelCase			)

      a     :			Union[str, Any]				=   model.init_cache(input_ids.shape[0]				,		__UpperCAmelCase			)
      a     :			Any				=   jnp.ones((input_ids.shape[0], max_decoder_length)				,		dtype='i4'			)

      a     :			Optional[Any]				=   jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1			)[None, :]				,		(input_ids.shape[0], input_ids.shape[-1] - 1)			)
      a     :			List[Any]				=   model(
          input_ids[:, :-1]				,		attention_mask=__UpperCAmelCase				,		past_key_values=__UpperCAmelCase				,		position_ids=__UpperCAmelCase				,		)

      a     :			Optional[int]				=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,		dtype='i4'			)
      a     :			List[Any]				=   model(
          input_ids[:, -1:]				,		attention_mask=__UpperCAmelCase				,		past_key_values=outputs_cache.past_key_values				,		position_ids=__UpperCAmelCase				,		)

      a     :			Optional[Any]				=   model(__UpperCAmelCase			)

      a     :			Optional[Any]				=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])			)			)
      self.parent.assertTrue(diff < 1E-3				,		msg=F'''Max diff is {diff}'''			)





    def 				lowerCamelCase__      (							self	:					Dict				,		A	:					Any				,		A	:					Tuple				,		A	:					List[Any]				,		A	:					int			):


      '''simple docstring'''




      a     :			int				=   2_0
      a     :			int				=   model_class_name(__UpperCAmelCase			)

      a     :			List[Any]				=   jnp.concatenate(
          [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])			)]				,		axis=-1				,		)

      a     :			Optional[int]				=   model.init_cache(input_ids.shape[0]				,		__UpperCAmelCase			)
      a     :			Tuple				=   jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1			)[None, :]				,		(input_ids.shape[0], input_ids.shape[-1] - 1)			)

      a     :			List[str]				=   model(
          input_ids[:, :-1]				,		attention_mask=__UpperCAmelCase				,		past_key_values=__UpperCAmelCase				,		position_ids=__UpperCAmelCase				,		)
      a     :			Optional[Any]				=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,		dtype='i4'			)
      a     :			Dict				=   model(
          input_ids[:, -1:]				,		past_key_values=outputs_cache.past_key_values				,		attention_mask=__UpperCAmelCase				,		position_ids=__UpperCAmelCase				,		)

      a     :			int				=   model(__UpperCAmelCase				,		attention_mask=__UpperCAmelCase			)
      a     :			Optional[Any]				=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])			)			)
      self.parent.assertTrue(diff < 1E-3				,		msg=F'''Max diff is {diff}'''			)






@require_sentencepiece
@require_flax
class 						snake_case  (  snake_case__   , snake_case__   , unittest.TestCase       ):
    __magic_name__         =     (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
    __magic_name__         =     (FlaxXGLMForCausalLM,) if is_flax_available() else ()





    def 				lowerCamelCase__      (							self	:					Tuple			):


      '''simple docstring'''




      a     :			Any				=   FlaxXGLMModelTester(self			)





    def 				lowerCamelCase__      (							self	:					Any			):


      '''simple docstring'''




      for model_class_name in self.all_model_classes:
        a,     a,     a     :			Union[str, Any]				=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward(__UpperCAmelCase				,		__UpperCAmelCase				,		__UpperCAmelCase				,		__UpperCAmelCase			)





    def 				lowerCamelCase__      (							self	:					Optional[int]			):


      '''simple docstring'''




      for model_class_name in self.all_model_classes:
        a,     a,     a     :			Optional[Any]				=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward_with_attn_mask(
            __UpperCAmelCase				,		__UpperCAmelCase				,		__UpperCAmelCase				,		__UpperCAmelCase			)





    @slow
    def 				lowerCamelCase__      (							self	:					Optional[int]			):


      '''simple docstring'''




      a     :			Optional[Any]				=   XGLMTokenizer.from_pretrained('XGLM'				,		padding_side='left'			)
      a     :			Dict				=   tokenizer(['Hello this is a long string', 'Hey']				,		return_tensors='np'				,		padding=__UpperCAmelCase				,		truncation=__UpperCAmelCase			)

      a     :			Union[str, Any]				=   FlaxXGLMForCausalLM.from_pretrained('facebook/xglm-564M'			)
      a     :			str				=   1
      a     :			Optional[int]				=   False

      a     :			str				=   jax.jit(model.generate			)

      a     :			List[str]				=   jit_generate(inputs['input_ids']				,		attention_mask=inputs['attention_mask']			).sequences

      a     :			List[Any]				=   tokenizer.batch_decode(__UpperCAmelCase				,		skip_special_tokens=__UpperCAmelCase			)

      a     :			Union[str, Any]				=   [
          'Hello this is a long string of questions, but I\'m not sure if I\'m',
          'Hey, I\'m a newbie to the forum and I\'',
      ]

      self.assertListEqual(__UpperCAmelCase				,		__UpperCAmelCase			)





    @is_pt_flax_cross_test
    def 				lowerCamelCase__      (							self	:					Tuple			):


      '''simple docstring'''




      a,     a     :			Tuple				=   self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__			):
          # prepare inputs
          a     :			Optional[int]				=   self._prepare_for_class(__UpperCAmelCase				,		__UpperCAmelCase			)
          a     :			List[str]				=   {k: torch.tensor(v.tolist()			) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          a     :			str				=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          a     :			Any				=   getattr(__UpperCAmelCase				,		__UpperCAmelCase			)

          a,     a     :			Any				=   pt_inputs['input_ids'].shape
          a     :			Union[str, Any]				=   np.random.randint(0				,		seq_length - 1				,		size=(batch_size,)			)
          for batch_idx, start_index in enumerate(__UpperCAmelCase			):
            a     :			List[Any]				=   0
            a     :			int				=   1
            a     :			int				=   0
            a     :			List[str]				=   1
          a     :			Tuple				=   pt_model_class(__UpperCAmelCase			).eval()
          # Flax models don't use the `use_cache` option and cache is not returned as a default.
          # So we disable `use_cache` here for PyTorch model.
          a     :			Optional[int]				=   False
          a     :			Tuple				=   model_class(__UpperCAmelCase				,		dtype=jnp.floataa			)
          a     :			Union[str, Any]				=   convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,		__UpperCAmelCase			)
          a     :			Any				=   fx_state

          with torch.no_grad():
            a     :			List[Any]				=   pt_model(**__UpperCAmelCase			).to_tuple()

          a     :			Dict				=   fx_model(**__UpperCAmelCase			).to_tuple()
          self.assertEqual(len(__UpperCAmelCase			)				,		len(__UpperCAmelCase			)				,		'Output lengths differ between Flax and PyTorch'			)
          for fx_output, pt_output in zip(__UpperCAmelCase				,		__UpperCAmelCase			):
            self.assert_almost_equals(fx_output[:, -1]				,		pt_output[:, -1].numpy()				,		4E-2			)

          with tempfile.TemporaryDirectory() as tmpdirname:
            pt_model.save_pretrained(__UpperCAmelCase			)
            a     :			Optional[int]				=   model_class.from_pretrained(__UpperCAmelCase				,		from_pt=__UpperCAmelCase			)

          a     :			Tuple				=   fx_model_loaded(**__UpperCAmelCase			).to_tuple()
          self.assertEqual(
              len(__UpperCAmelCase			)				,		len(__UpperCAmelCase			)				,		'Output lengths differ between Flax and PyTorch'			)
          for fx_output_loaded, pt_output in zip(__UpperCAmelCase				,		__UpperCAmelCase			):
            self.assert_almost_equals(fx_output_loaded[:, -1]				,		pt_output[:, -1].numpy()				,		4E-2			)





    @is_pt_flax_cross_test
    def 				lowerCamelCase__      (							self	:					Tuple			):


      '''simple docstring'''




      a,     a     :			Dict				=   self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__			):
          # prepare inputs
          a     :			str				=   self._prepare_for_class(__UpperCAmelCase				,		__UpperCAmelCase			)
          a     :			Tuple				=   {k: torch.tensor(v.tolist()			) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          a     :			int				=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          a     :			Optional[Any]				=   getattr(__UpperCAmelCase				,		__UpperCAmelCase			)

          a     :			Union[str, Any]				=   pt_model_class(__UpperCAmelCase			).eval()
          a     :			Dict				=   False
          a     :			str				=   model_class(__UpperCAmelCase				,		dtype=jnp.floataa			)

          a     :			List[Any]				=   load_flax_weights_in_pytorch_model(__UpperCAmelCase				,		fx_model.params			)
          a,     a     :			Union[str, Any]				=   pt_inputs['input_ids'].shape
          a     :			Tuple				=   np.random.randint(0				,		seq_length - 1				,		size=(batch_size,)			)
          for batch_idx, start_index in enumerate(__UpperCAmelCase			):
            a     :			Optional[int]				=   0
            a     :			Any				=   1
            a     :			List[str]				=   0
            a     :			Any				=   1

          # make sure weights are tied in PyTorch
          pt_model.tie_weights()

          with torch.no_grad():
            a     :			List[Any]				=   pt_model(**__UpperCAmelCase			).to_tuple()

          a     :			Tuple				=   fx_model(**__UpperCAmelCase			).to_tuple()
          self.assertEqual(len(__UpperCAmelCase			)				,		len(__UpperCAmelCase			)				,		'Output lengths differ between Flax and PyTorch'			)
          for fx_output, pt_output in zip(__UpperCAmelCase				,		__UpperCAmelCase			):
            self.assert_almost_equals(fx_output[:, -1]				,		pt_output[:, -1].numpy()				,		4E-2			)

          with tempfile.TemporaryDirectory() as tmpdirname:
            fx_model.save_pretrained(__UpperCAmelCase			)
            a     :			Union[str, Any]				=   pt_model_class.from_pretrained(__UpperCAmelCase				,		from_flax=__UpperCAmelCase			)

          with torch.no_grad():
            a     :			List[str]				=   pt_model_loaded(**__UpperCAmelCase			).to_tuple()

          self.assertEqual(
              len(__UpperCAmelCase			)				,		len(__UpperCAmelCase			)				,		'Output lengths differ between Flax and PyTorch'			)
          for fx_output, pt_output in zip(__UpperCAmelCase				,		__UpperCAmelCase			):
            self.assert_almost_equals(fx_output[:, -1]				,		pt_output[:, -1].numpy()				,		4E-2			)





    @slow
    def 				lowerCamelCase__      (							self	:					Union[str, Any]			):


      '''simple docstring'''




      for model_class_name in self.all_model_classes:
        a     :			Dict				=   model_class_name.from_pretrained('facebook/xglm-564M'			)
        a     :			Any				=   model(np.ones((1, 1)			)			)
        self.assertIsNotNone(__UpperCAmelCase			)









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)




