






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


def      lowerCamelCase__	(   lowercase  ,   lowercase  ,   lowercase  ,   lowercase  ,   lowercase  ,   lowercase  ,   lowercase=None ):

  """simple docstring"""






  SCREAMING_SNAKE_CASE				:					int		  =			tqdm(lowercase )
  for i, batch in enumerate(lowercase ):
    with tf.GradientTape() as tape:
      SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE				:					str		  =			batch
      SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			model(lowercase )
      SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			loss(lowercase  ,   outputs.logits )

    if SDP_ENABLED:
      SCREAMING_SNAKE_CASE				:					Any		  =			sdp.DistributedGradientTape(lowercase  ,   sparse_as_dense=lowercase )

    SCREAMING_SNAKE_CASE				:					Tuple		  =			tape.gradient(lowercase  ,   model.trainable_variables )
    opt.apply_gradients(zip(lowercase  ,   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

  SCREAMING_SNAKE_CASE				:					Optional[int]		  =			{"loss": loss_value.numpy()}
  return train_results


def      lowerCamelCase__	(   lowercase  ,   lowercase  ,   lowercase ):

  """simple docstring"""






  SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			load_dataset("imdb"  ,   split=["train", "test"] )

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

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

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

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

  if SDP_ENABLED:
    SCREAMING_SNAKE_CASE				:					int		  =			tf_train_dataset.shard(sdp.size()  ,   sdp.rank() )
    SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			tf_test_dataset.shard(sdp.size()  ,   sdp.rank() )
  SCREAMING_SNAKE_CASE				:					List[str]		  =			tf_train_dataset.batch(lowercase  ,   drop_remainder=lowercase )
  SCREAMING_SNAKE_CASE				:					str		  =			tf_test_dataset.batch(lowercase  ,   drop_remainder=lowercase )

  return tf_train_dataset, tf_test_dataset


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

   snake_case ,      snake_case			 =						parser.parse_known_args()

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

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

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

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

      snake_case			 =						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)














import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
       SpeechaTextaDecoder,
       SpeechaTextaForCausalLM,
   )


@require_torch
class        SCREAMING_SNAKE_CASE	:
     '''simple docstring'''


     def __init__( self				:		List[str]    ,       UpperCAmelCase_				:		Any    ,       UpperCAmelCase_				:		List[Any]=99    ,       UpperCAmelCase_				:		Tuple=13    ,       UpperCAmelCase_				:		Dict=16    ,       UpperCAmelCase_				:		int=7    ,       UpperCAmelCase_				:		Optional[Any]=True    ,       UpperCAmelCase_				:		List[Any]=True    ,       UpperCAmelCase_				:		Any=True    ,       UpperCAmelCase_				:		List[str]=False    ,       UpperCAmelCase_				:		Tuple=True    ,       UpperCAmelCase_				:		Tuple=2    ,       UpperCAmelCase_				:		List[str]=32    ,       UpperCAmelCase_				:		Union[str, Any]=4    ,       UpperCAmelCase_				:		Any=4    ,       UpperCAmelCase_				:		Optional[int]=30    ,       UpperCAmelCase_				:		Union[str, Any]=0    ,       UpperCAmelCase_				:		Union[str, Any]=1    ,       UpperCAmelCase_				:		Any=2    ,       UpperCAmelCase_				:		Dict=None    ,       ):
       SCREAMING_SNAKE_CASE				:					Any		  =			parent
       SCREAMING_SNAKE_CASE				:					Tuple		  =			batch_size
       SCREAMING_SNAKE_CASE				:					str		  =			decoder_seq_length
       # For common tests
       SCREAMING_SNAKE_CASE				:					Tuple		  =			self.decoder_seq_length
       SCREAMING_SNAKE_CASE				:					str		  =			is_training
       SCREAMING_SNAKE_CASE				:					Dict		  =			use_attention_mask
       SCREAMING_SNAKE_CASE				:					int		  =			use_labels

       SCREAMING_SNAKE_CASE				:					Tuple		  =			vocab_size
       SCREAMING_SNAKE_CASE				:					int		  =			d_model
       SCREAMING_SNAKE_CASE				:					List[str]		  =			d_model
       SCREAMING_SNAKE_CASE				:					int		  =			decoder_layers
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			decoder_layers
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			decoder_ffn_dim
       SCREAMING_SNAKE_CASE				:					Tuple		  =			decoder_attention_heads
       SCREAMING_SNAKE_CASE				:					Dict		  =			decoder_attention_heads
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			eos_token_id
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			bos_token_id
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			pad_token_id
       SCREAMING_SNAKE_CASE				:					List[str]		  =			decoder_start_token_id
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			use_cache
       SCREAMING_SNAKE_CASE				:					Tuple		  =			max_position_embeddings

       SCREAMING_SNAKE_CASE				:					Any		  =			None
       SCREAMING_SNAKE_CASE				:					Tuple		  =			decoder_seq_length
       SCREAMING_SNAKE_CASE				:					int		  =			2
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			1


     def  _A				( self				:		Optional[int]					):
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			ids_tensor([self.batch_size, self.decoder_seq_length]    ,       self.vocab_size					)

       SCREAMING_SNAKE_CASE				:					str		  =			None
       if self.use_attention_mask:
         SCREAMING_SNAKE_CASE				:					Tuple		  =			ids_tensor([self.batch_size, self.decoder_seq_length]    ,       vocab_size=2					)

       SCREAMING_SNAKE_CASE				:					str		  =			None
       if self.use_labels:
         SCREAMING_SNAKE_CASE				:					List[str]		  =			ids_tensor([self.batch_size, self.decoder_seq_length]    ,       self.vocab_size					)

       SCREAMING_SNAKE_CASE				:					str		  =			SpeechaTextaConfig(
           vocab_size=self.vocab_size    ,       d_model=self.d_model    ,       decoder_layers=self.decoder_layers    ,       decoder_ffn_dim=self.decoder_ffn_dim    ,       decoder_attention_heads=self.decoder_attention_heads    ,       eos_token_id=self.eos_token_id    ,       bos_token_id=self.bos_token_id    ,       use_cache=self.use_cache    ,       pad_token_id=self.pad_token_id    ,       decoder_start_token_id=self.decoder_start_token_id    ,       max_position_embeddings=self.max_position_embeddings    ,       )

       return (
           config,
           input_ids,
           attention_mask,
           lm_labels,
       )


     def  _A				( self				:		Optional[Any]    ,       UpperCAmelCase_				:		int    ,       UpperCAmelCase_				:		Dict    ,       UpperCAmelCase_				:		int    ,       UpperCAmelCase_				:		Any    ,       ):
       SCREAMING_SNAKE_CASE				:					str		  =			True
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			SpeechaTextaDecoder(config=UpperCAmelCase_					).to(UpperCAmelCase_					).eval()
       SCREAMING_SNAKE_CASE				:					Tuple		  =			input_ids[:2]

       input_ids[input_ids == 0] += 1
       # first forward pass
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			model(UpperCAmelCase_    ,       use_cache=UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					List[str]		  =			model(UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			model(UpperCAmelCase_    ,       use_cache=UpperCAmelCase_					)

       self.parent.assertTrue(len(UpperCAmelCase_					) == len(UpperCAmelCase_					)					)
       self.parent.assertTrue(len(UpperCAmelCase_					) == len(UpperCAmelCase_					) + 1					)

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			outputs["past_key_values"]

       # create hypothetical next token and extent to next_input_ids
       SCREAMING_SNAKE_CASE				:					str		  =			ids_tensor((2, 1)    ,       config.vocab_size - 1					) + 1

       # append to next input_ids and
       SCREAMING_SNAKE_CASE				:					str		  =			torch.cat([input_ids, next_tokens]    ,       dim=-1					)

       print(UpperCAmelCase_					)

       SCREAMING_SNAKE_CASE				:					List[str]		  =			model(UpperCAmelCase_					)["last_hidden_state"]
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			model(UpperCAmelCase_    ,       past_key_values=UpperCAmelCase_					)["last_hidden_state"]

       # select random slice
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			ids_tensor((1,)    ,       output_from_past.shape[-1]					).item()
       SCREAMING_SNAKE_CASE				:					str		  =			output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
       SCREAMING_SNAKE_CASE				:					int		  =			output_from_past[:, 0, random_slice_idx].detach()

       # test that outputs are equal for slice
       assert torch.allclose(UpperCAmelCase_    ,       UpperCAmelCase_    ,       atol=1E-3					)






     def  _A				( self				:		str					):
       SCREAMING_SNAKE_CASE				:					int		  =			self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE
       )  ,     (
           SCREAMING_SNAKE_CASE
       )  ,     (
           SCREAMING_SNAKE_CASE
       )  ,     (
           SCREAMING_SNAKE_CASE
       )  ,     
       )				:					Optional[Any]		  =			config_and_inputs

       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			{
           "input_ids": input_ids,
           "attention_mask": attention_mask,
       }
       return config, inputs_dict





@require_torch
class        SCREAMING_SNAKE_CASE	(	lowerCAmelCase					,						lowerCAmelCase					,						lowerCAmelCase					,						unittest.TestCase    ):
     '''simple docstring'''


     UpperCamelCase_							:						int								=		(SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
     UpperCamelCase_							:						Tuple								=		(SpeechaTextaForCausalLM,) if is_torch_available() else ()
     UpperCamelCase_							:						int								=		{'''text-generation''': SpeechaTextaForCausalLM} if is_torch_available() else {}
     UpperCamelCase_							:						Tuple								=		True
     UpperCamelCase_							:						Union[str, Any]								=		False


     def  _A				( self				:		Tuple    ,       ):
       SCREAMING_SNAKE_CASE				:					Any		  =			SpeechaTextaStandaloneDecoderModelTester(self    ,       is_training=UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			ConfigTester(self    ,       config_class=UpperCAmelCase_					)


     def  _A				( self				:		Optional[Any]					):
       pass


     def  _A				( self				:		str					):
       pass


     def  _A				( self				:		Dict					):
       pass


     def  _A				( self				:		Optional[Any]					):
       self.config_tester.run_common_tests()


     def  _A				( self				:		Optional[int]					):
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_decoder_model_past(*UpperCAmelCase_					)






     def  _A				( self				:		Optional[int]					):
       return






