
from __future__ import annotations

import unittest

import numpy as np

from transformers import OPTConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import GPTaTokenizer, TFOPTForCausalLM, TFOPTModel


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__ ,		UpperCamelCase__=None ,		UpperCamelCase__=None		):
      if attention_mask is None:
            UpperCamelCase__    :   Union[str, Any]					=				tf.cast(tf.math.not_equal(UpperCamelCase__ ,		config.pad_token_id		) ,		tf.inta		)
      return {"input_ids": input_ids, "attention_mask": attention_mask}




@require_tf
class 				_lowerCamelCase :





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  OPTConfig
     SCREAMING_SNAKE_CASE_		   =  {}
     SCREAMING_SNAKE_CASE_		   =  '''gelu'''






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=1_3						,						__SCREAMING_SNAKE_CASE=7						,						__SCREAMING_SNAKE_CASE=True						,						__SCREAMING_SNAKE_CASE=False						,						__SCREAMING_SNAKE_CASE=9_9						,						__SCREAMING_SNAKE_CASE=1_6						,						__SCREAMING_SNAKE_CASE=2						,						__SCREAMING_SNAKE_CASE=4						,						__SCREAMING_SNAKE_CASE=4						,						__SCREAMING_SNAKE_CASE="gelu"						,						__SCREAMING_SNAKE_CASE=0.1						,						__SCREAMING_SNAKE_CASE=0.1						,						__SCREAMING_SNAKE_CASE=2_0						,						__SCREAMING_SNAKE_CASE=2						,						__SCREAMING_SNAKE_CASE=1						,						__SCREAMING_SNAKE_CASE=0						,						__SCREAMING_SNAKE_CASE=1_6						,						__SCREAMING_SNAKE_CASE=1_6						,						)  ->   int:




           """simple docstring"""







           UpperCamelCase__    :   Tuple					=				parent
           UpperCamelCase__    :   Tuple					=				batch_size
           UpperCamelCase__    :   List[Any]					=				seq_length
           UpperCamelCase__    :   Dict					=				is_training
           UpperCamelCase__    :   List[str]					=				use_labels
           UpperCamelCase__    :   Optional[Any]					=				vocab_size
           UpperCamelCase__    :   str					=				hidden_size
           UpperCamelCase__    :   List[Any]					=				num_hidden_layers
           UpperCamelCase__    :   Dict					=				num_attention_heads
           UpperCamelCase__    :   Any					=				intermediate_size
           UpperCamelCase__    :   List[Any]					=				hidden_act
           UpperCamelCase__    :   List[Any]					=				hidden_dropout_prob
           UpperCamelCase__    :   Union[str, Any]					=				attention_probs_dropout_prob
           UpperCamelCase__    :   Dict					=				max_position_embeddings
           UpperCamelCase__    :   List[Any]					=				eos_token_id
           UpperCamelCase__    :   str					=				pad_token_id
           UpperCamelCase__    :   int					=				bos_token_id
           UpperCamelCase__    :   Tuple					=				embed_dim
           UpperCamelCase__    :   Dict					=				word_embed_proj_dim
           UpperCamelCase__    :   str					=				False






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Tuple:




           """simple docstring"""







           UpperCamelCase__    :   List[str]					=				ids_tensor([self.batch_size, self.seq_length - 1]						,						self.vocab_size  )
           UpperCamelCase__    :   Dict					=				tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size  )						,						1  )
           UpperCamelCase__    :   Any					=				tf.concat([input_ids, eos_tensor]						,						axis=1  )

           UpperCamelCase__    :   Dict					=				self.config_cls(
               vocab_size=self.vocab_size						,						hidden_size=self.hidden_size						,						num_hidden_layers=self.num_hidden_layers						,						num_attention_heads=self.num_attention_heads						,						ffn_dim=self.intermediate_size						,						dropout=self.hidden_dropout_prob						,						attention_dropout=self.attention_probs_dropout_prob						,						max_position_embeddings=self.max_position_embeddings						,						eos_token_id=self.eos_token_id						,						bos_token_id=self.bos_token_id						,						pad_token_id=self.pad_token_id						,						embed_dim=self.embed_dim						,						word_embed_proj_dim=self.word_embed_proj_dim						,						is_encoder_decoder=__SCREAMING_SNAKE_CASE						,						**self.config_updates						,						)
           UpperCamelCase__    :   Any					=				prepare_opt_inputs_dict(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )
           return config, inputs_dict






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )  ->   Optional[int]:




           """simple docstring"""







           UpperCamelCase__    :   str					=				TFOPTModel(config=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Any					=				inputs_dict['''input_ids''']

           UpperCamelCase__    :   List[Any]					=				input_ids[:1, :]
           UpperCamelCase__    :   Optional[Any]					=				inputs_dict['''attention_mask'''][:1, :]
           UpperCamelCase__    :   Tuple					=				1

           # first forward pass
           UpperCamelCase__    :   Union[str, Any]					=				model(__SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						use_cache=__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__				,UpperCamelCase__    :   Optional[Any]					=				outputs.to_tuple()

           # create hypothetical next token and extent to next_input_ids
           UpperCamelCase__    :   int					=				ids_tensor((self.batch_size, 3)						,						config.vocab_size  )
           UpperCamelCase__    :   Tuple					=				tf.cast(ids_tensor((self.batch_size, 3)						,						2  )						,						tf.inta  )

           # append to next input_ids and
           UpperCamelCase__    :   Optional[Any]					=				tf.concat([input_ids, next_tokens]						,						axis=-1  )
           UpperCamelCase__    :   List[Any]					=				tf.concat([attention_mask, next_attn_mask]						,						axis=-1  )

           UpperCamelCase__    :   List[str]					=				model(__SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE  )[0]
           UpperCamelCase__    :   Optional[Any]					=				model(__SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						past_key_values=__SCREAMING_SNAKE_CASE  )[0]

           self.parent.assertEqual(next_tokens.shape[1]						,						output_from_past.shape[1]  )

           # select random slice
           UpperCamelCase__    :   int					=				int(ids_tensor((1,)						,						output_from_past.shape[-1]  )  )
           UpperCamelCase__    :   Union[str, Any]					=				output_from_no_past[:, -3:, random_slice_idx]
           UpperCamelCase__    :   List[str]					=				output_from_past[:, :, random_slice_idx]

           # test that outputs are equal for slice
           tf.debugging.assert_near(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						rtol=1e-3  )




@require_tf
class 				_lowerCamelCase (      UpperCamelCase_	,							UpperCamelCase_	,							unittest.TestCase ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
     SCREAMING_SNAKE_CASE_		   =  (TFOPTForCausalLM,) if is_tf_available() else ()
     SCREAMING_SNAKE_CASE_		   =  (
         {'''feature-extraction''': TFOPTModel, '''text-generation''': TFOPTForCausalLM} if is_tf_available() else {}
     )
     SCREAMING_SNAKE_CASE_		   =  False
     SCREAMING_SNAKE_CASE_		   =  False
     SCREAMING_SNAKE_CASE_		   =  False
     SCREAMING_SNAKE_CASE_		   =  10






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           UpperCamelCase__    :   Optional[int]					=				TFOPTModelTester(self  )
           UpperCamelCase__    :   Any					=				ConfigTester(self						,						config_class=__SCREAMING_SNAKE_CASE  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







           self.config_tester.run_common_tests()






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   int:




           """simple docstring"""







           UpperCamelCase__    :   Union[str, Any]					=				self.model_tester.prepare_config_and_inputs_for_common()
           self.model_tester.check_decoder_model_past_large_inputs(*__SCREAMING_SNAKE_CASE  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




           """simple docstring"""







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

           def _get_word_embedding_weight(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  ):
                 if hasattr(__SCREAMING_SNAKE_CASE						,						'''weight'''  ):
                       return embedding_layer.weight
                 else:
                       # Here we build the word embeddings weights if not exists.
                       # And then we retry to get the attribute once built.
                       model.build()
                       if hasattr(__SCREAMING_SNAKE_CASE						,						'''weight'''  ):
                             return embedding_layer.weight
                       else:
                             return None

           for model_class in self.all_model_classes:
                 for size in [config.vocab_size - 1_0, config.vocab_size + 1_0]:
                       # build the embeddings
                       UpperCamelCase__    :   str					=				model_class(config=__SCREAMING_SNAKE_CASE  )
                       UpperCamelCase__    :   List[str]					=				_get_word_embedding_weight(__SCREAMING_SNAKE_CASE						,						model.get_input_embeddings()  )
                       UpperCamelCase__    :   List[Any]					=				_get_word_embedding_weight(__SCREAMING_SNAKE_CASE						,						model.get_output_embeddings()  )

                       # reshape the embeddings
                       model.resize_token_embeddings(__SCREAMING_SNAKE_CASE  )
                       UpperCamelCase__    :   Optional[Any]					=				_get_word_embedding_weight(__SCREAMING_SNAKE_CASE						,						model.get_input_embeddings()  )
                       UpperCamelCase__    :   Dict					=				_get_word_embedding_weight(__SCREAMING_SNAKE_CASE						,						model.get_output_embeddings()  )

                       # check that the resized embeddings size matches the desired size.
                       UpperCamelCase__    :   List[str]					=				size if size is not None else config.vocab_size

                       self.assertEqual(new_input_embeddings.shape[0]						,						__SCREAMING_SNAKE_CASE  )

                       # check that weights remain the same after resizing
                       UpperCamelCase__    :   str					=				True
                       for pa, pa in zip(old_input_embeddings.value()						,						new_input_embeddings.value()  ):
                             if tf.math.reduce_sum(tf.math.abs(pa - pa  )  ) > 0:
                                   UpperCamelCase__    :   Any					=				False
                       self.assertTrue(__SCREAMING_SNAKE_CASE  )

                       if old_output_embeddings is not None and new_output_embeddings is not None:
                             self.assertEqual(new_output_embeddings.shape[0]						,						__SCREAMING_SNAKE_CASE  )

                             UpperCamelCase__    :   Optional[int]					=				True
                             for pa, pa in zip(old_output_embeddings.value()						,						new_output_embeddings.value()  ):
                                   if tf.math.reduce_sum(tf.math.abs(pa - pa  )  ) > 0:
                                         UpperCamelCase__    :   List[Any]					=				False
                             self.assertTrue(__SCREAMING_SNAKE_CASE  )


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      return tf.constant(UpperCamelCase__ ,		dtype=tf.intaa		)




@require_tf
class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  99






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[int]:




           """simple docstring"""







           UpperCamelCase__    :   Union[str, Any]					=				tf.ones((4, 1)						,						dtype=tf.intaa  ) * 2
           UpperCamelCase__    :   Union[str, Any]					=				tf.concat([ids_tensor((4, 6)						,						self.vocab_size - 3  ) + 3, eos_column_vector]						,						axis=1  )
           UpperCamelCase__    :   List[str]					=				input_ids.shape[0]
           UpperCamelCase__    :   Union[str, Any]					=				OPTConfig(
               vocab_size=self.vocab_size						,						hidden_size=2_4						,						num_hidden_layers=2						,						num_attention_heads=2						,						ffn_dim=3_2						,						max_position_embeddings=4_8						,						eos_token_id=2						,						pad_token_id=1						,						bos_token_id=0						,						)
           return config, input_ids, batch_size




@require_sentencepiece
@require_tf
class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     @slow
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Union[str, Any]:




           """simple docstring"""







           UpperCamelCase__    :   str					=				TFOPTModel.from_pretrained('''facebook/opt-350m'''  )
           UpperCamelCase__    :   Optional[int]					=				_long_tensor([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]  )
           UpperCamelCase__    :   Any					=				tf.not_equal(__SCREAMING_SNAKE_CASE						,						model.config.pad_token_id  )
           with tf.GradientTape():
                 UpperCamelCase__    :   Optional[Any]					=				model(input_ids=__SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE  ).last_hidden_state
           UpperCamelCase__    :   Tuple					=				(1, 1_1, 5_1_2)
           self.assertEqual(output.shape						,						__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Dict					=				tf.constant(
               [[-0.2873, -1.9218, -0.3033], [-1.2710, -0.1338, -0.1902], [0.4095, 0.1214, -1.3121]]  )
           self.assertTrue(np.allclose(output[:, :3, :3]						,						__SCREAMING_SNAKE_CASE						,						atol=4e-3  )  )

           UpperCamelCase__    :   Any					=				tf.function(__SCREAMING_SNAKE_CASE						,						jit_compile=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   str					=				xla_generate(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )[0]
           self.assertTrue(np.allclose(output[:, :3, :3]						,						__SCREAMING_SNAKE_CASE						,						atol=4e-2  )  )




@require_tf
@slow
class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   int:




           """simple docstring"""







           super().setUp()
           UpperCamelCase__    :   Optional[int]					=				'''facebook/opt-350m'''






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           UpperCamelCase__    :   List[str]					=				TFOPTForCausalLM.from_pretrained(self.path_model  )
           UpperCamelCase__    :   Any					=				GPTaTokenizer.from_pretrained(self.path_model  )

           UpperCamelCase__    :   Dict					=				[
               '''Today is a beautiful day and I want to''',
               '''In the city of''',
               '''Paris is the capital of France and''',
               '''Computers and mobile phones have taken''',
           ]
           # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
           UpperCamelCase__    :   Tuple					=				tokenizer(__SCREAMING_SNAKE_CASE						,						return_tensors='''tf'''						,						padding=__SCREAMING_SNAKE_CASE						,						add_special_tokens=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Optional[int]					=				tf.math.reduce_mean(model(inputs.input_ids						,						attention_mask=inputs.attention_mask  )[0]						,						axis=-1  )
           UpperCamelCase__    :   Optional[Any]					=				tf.constant(
               [
                   [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
                   [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
                   [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
                   [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
               ]  )
           self.assertTrue(np.allclose(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						atol=1e-4  )  )

           UpperCamelCase__    :   Optional[int]					=				tf.function(__SCREAMING_SNAKE_CASE						,						jit_compile=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Any					=				tf.math.reduce_mean(xla_generate(inputs.input_ids						,						attention_mask=inputs.attention_mask  )[0]						,						axis=-1  )
           self.assertTrue(np.allclose(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						atol=1e-4  )  )




@require_tf
@slow
class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     @property
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           return [
               "Today is a beautiful day and I want",
               "In the city of",
               "Paris is the capital of France and",
               "Computers and mobile phones have taken",
           ]






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




           """simple docstring"""







           UpperCamelCase__    :   int					=				'''facebook/opt-125m'''

           UpperCamelCase__    :   Optional[int]					=				[
               '''Today is a beautiful day and I want to''',
               '''In the city of New York, the city''',
               '''Paris is the capital of France and the capital''',
               '''Computers and mobile phones have taken over the''',
           ]

           UpperCamelCase__    :   List[str]					=				[]
           UpperCamelCase__    :   List[Any]					=				GPTaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   List[Any]					=				TFOPTForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE  )

           for prompt in self.prompts:
                 UpperCamelCase__    :   Union[str, Any]					=				tokenizer(__SCREAMING_SNAKE_CASE						,						return_tensors='''tf'''  ).input_ids

                 UpperCamelCase__    :   Union[str, Any]					=				model.generate(__SCREAMING_SNAKE_CASE						,						max_length=1_0  )

                 UpperCamelCase__    :   Optional[int]					=				tokenizer.batch_decode(__SCREAMING_SNAKE_CASE						,						skip_special_tokens=__SCREAMING_SNAKE_CASE  )
                 predicted_outputs += generated_string

           self.assertListEqual(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




           """simple docstring"""







           UpperCamelCase__    :   Optional[int]					=				'''facebook/opt-350m'''

           UpperCamelCase__    :   str					=				GPTaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Tuple					=				TFOPTForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__    :   Optional[Any]					=				'''left'''

           # use different length sentences to test batching
           UpperCamelCase__    :   Dict					=				[
               '''Hello, my dog is a little''',
               '''Today, I''',
           ]

           UpperCamelCase__    :   str					=				tokenizer(__SCREAMING_SNAKE_CASE						,						return_tensors='''tf'''						,						padding=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Dict					=				inputs['''input_ids''']

           UpperCamelCase__    :   List[Any]					=				model.generate(input_ids=__SCREAMING_SNAKE_CASE						,						attention_mask=inputs['''attention_mask''']  )

           UpperCamelCase__    :   Any					=				tokenizer(sentences[0]						,						return_tensors='''tf'''  ).input_ids
           UpperCamelCase__    :   str					=				model.generate(input_ids=__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__    :   Tuple					=				inputs_non_padded.shape[-1] - tf.math.reduce_sum(
               tf.cast(inputs['''attention_mask'''][-1]						,						tf.intaa  )  )
           UpperCamelCase__    :   List[str]					=				tokenizer(sentences[1]						,						return_tensors='''tf'''  ).input_ids
           UpperCamelCase__    :   Optional[int]					=				model.generate(input_ids=__SCREAMING_SNAKE_CASE						,						max_length=model.config.max_length - num_paddings  )

           UpperCamelCase__    :   Any					=				tokenizer.batch_decode(__SCREAMING_SNAKE_CASE						,						skip_special_tokens=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   str					=				tokenizer.decode(output_non_padded[0]						,						skip_special_tokens=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Dict					=				tokenizer.decode(output_padded[0]						,						skip_special_tokens=__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__    :   Dict					=				[
               '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
               '''Today, I was in the middle of a conversation with a friend about the''',
           ]
           self.assertListEqual(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )
           self.assertListEqual(__SCREAMING_SNAKE_CASE						,						[non_padded_sentence, padded_sentence]  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Tuple:




           """simple docstring"""







           UpperCamelCase__    :   Union[str, Any]					=				'''facebook/opt-350m'''

           UpperCamelCase__    :   int					=				[
               '''Today is a beautiful day and I want to''',
               '''In the city of San Francisco, the city''',
               '''Paris is the capital of France and the capital''',
               '''Computers and mobile phones have taken over the''',
           ]

           UpperCamelCase__    :   Optional[Any]					=				[]
           UpperCamelCase__    :   Tuple					=				GPTaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Dict					=				TFOPTForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE  )

           for prompt in self.prompts:
                 UpperCamelCase__    :   List[str]					=				tokenizer(__SCREAMING_SNAKE_CASE						,						return_tensors='''tf'''  ).input_ids

                 UpperCamelCase__    :   str					=				model.generate(__SCREAMING_SNAKE_CASE						,						max_length=1_0  )

                 UpperCamelCase__    :   Union[str, Any]					=				tokenizer.batch_decode(__SCREAMING_SNAKE_CASE						,						skip_special_tokens=__SCREAMING_SNAKE_CASE  )
                 predicted_outputs += generated_string

           self.assertListEqual(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )




import argparse
import shlex

import runhouse as rh


if __name__ == "__main__":
    # Refer to https://runhouse-docs.readthedocs-hosted.com/en/latest/api/python/cluster.html#hardware-setup for cloud access
    # setup instructions, if using on-demand hardware

    # If user passes --user <user> --host <host> --key_path <key_path> <example> <args>, fill them in as BYO cluster
    # If user passes --instance <instance> --provider <provider> <example> <args>, fill them in as on-demand cluster
    # Throw an error if user passes both BYO and on-demand cluster args
    # Otherwise, use default values
    lowerCamelCase      =argparse.ArgumentParser()
    parser.add_argument("--user", type=str, default="ubuntu")
    parser.add_argument("--host", type=str, default="localhost")
    parser.add_argument("--key_path", type=str, default=None)
    parser.add_argument("--instance", type=str, default="V100:1")
    parser.add_argument("--provider", type=str, default="cheapest")
    parser.add_argument("--use_spot", type=bool, default=False)
    parser.add_argument("--example", type=str, default="pytorch/text-generation/run_generation.py")
    lowerCamelCase   ,     lowerCamelCase      =parser.parse_known_args()
    if args.host != "localhost":
        if args.instance != "V100:1" or args.provider != "cheapest":
            raise ValueError("Cannot specify both BYO and on-demand cluster args")
        lowerCamelCase      =rh.cluster(
            name="rh-cluster", ips=[args.host], ssh_creds={"ssh_user": args.user, "ssh_private_key": args.key_path}
        )
    else:
        lowerCamelCase      =rh.cluster(
            name="rh-cluster", instance_type=args.instance, provider=args.provider, use_spot=args.use_spot
        )
    lowerCamelCase      =args.example.rsplit("/", 1)[0]

    # Set up remote environment
    cluster.install_packages(["pip:./"])  # Installs transformers from local source
    # Note transformers is copied into the home directory on the remote machine, so we can install from there
    cluster.run([F'''pip install -r transformers/examples/{example_dir}/requirements.txt'''])
    cluster.run(["pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117"])

    # Run example. You can bypass the CLI wrapper and paste your own code here.
    cluster.run([F'''python transformers/examples/{args.example} {" ".join(shlex.quote(arg) for arg in unknown)}'''])

    # Alternatively, we can just import and run a training function (especially if there's no wrapper CLI):
    # from my_script... import train
    # reqs = ['pip:./', 'torch', 'datasets', 'accelerate', 'evaluate', 'tqdm', 'scipy', 'scikit-learn', 'tensorboard']
    # launch_train_gpu = rh.function(fn=train,
    #                                system=gpu,
    #                                reqs=reqs,
    #                                name='train_bert_glue')
    #
    # We can pass in arguments just like we would to a function:
    # launch_train_gpu(num_epochs = 3, lr = 2e-5, seed = 42, batch_size = 16
    #                  stream_logs=True)


