'''simple docstring'''




import unittest

from transformers import LayoutLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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 import (
     LayoutLMForMaskedLM,
     LayoutLMForQuestionAnswering,
     LayoutLMForSequenceClassification,
     LayoutLMForTokenClassification,
     LayoutLMModel,
 )





class        SCREAMING_SNAKE_CASE   :
  """simple docstring"""







  def __init__(  self   :					Optional[Any]					,					__lowerCAmelCase   :					Dict					,					__lowerCAmelCase   :					Optional[Any]=13					,					__lowerCAmelCase   :					List[str]=7					,					__lowerCAmelCase   :					Tuple=True					,					__lowerCAmelCase   :					Dict=True					,					__lowerCAmelCase   :					Tuple=True					,					__lowerCAmelCase   :					List[str]=True					,					__lowerCAmelCase   :					int=99					,					__lowerCAmelCase   :					Optional[int]=32					,					__lowerCAmelCase   :					Union[str, Any]=5					,					__lowerCAmelCase   :					Dict=4					,					__lowerCAmelCase   :					Tuple=37					,					__lowerCAmelCase   :					int="gelu"					,					__lowerCAmelCase   :					Optional[Any]=0.1					,					__lowerCAmelCase   :					str=0.1					,					__lowerCAmelCase   :					List[Any]=512					,					__lowerCAmelCase   :					Tuple=16					,					__lowerCAmelCase   :					Dict=2					,					__lowerCAmelCase   :					str=0.0_2					,					__lowerCAmelCase   :					List[Any]=3					,					__lowerCAmelCase   :					Dict=4					,					__lowerCAmelCase   :					List[str]=None					,					__lowerCAmelCase   :					int=1000					,					):


        """simple docstring"""



        _lowerCAmelCase		   =   parent
        _lowerCAmelCase		   =   batch_size
        _lowerCAmelCase		   =   seq_length
        _lowerCAmelCase		   =   is_training
        _lowerCAmelCase		   =   use_input_mask
        _lowerCAmelCase		   =   use_token_type_ids
        _lowerCAmelCase		   =   use_labels
        _lowerCAmelCase		   =   vocab_size
        _lowerCAmelCase		   =   hidden_size
        _lowerCAmelCase		   =   num_hidden_layers
        _lowerCAmelCase		   =   num_attention_heads
        _lowerCAmelCase		   =   intermediate_size
        _lowerCAmelCase		   =   hidden_act
        _lowerCAmelCase		   =   hidden_dropout_prob
        _lowerCAmelCase		   =   attention_probs_dropout_prob
        _lowerCAmelCase		   =   max_position_embeddings
        _lowerCAmelCase		   =   type_vocab_size
        _lowerCAmelCase		   =   type_sequence_label_size
        _lowerCAmelCase		   =   initializer_range
        _lowerCAmelCase		   =   num_labels
        _lowerCAmelCase		   =   num_choices
        _lowerCAmelCase		   =   scope
        _lowerCAmelCase		   =   range_bbox



  def   a   (  self   :					Optional[int]					):


        """simple docstring"""



        _lowerCAmelCase		   =   ids_tensor([self.batch_size, self.seq_length]					,					self.vocab_size					)

        _lowerCAmelCase		   =   ids_tensor([self.batch_size, self.seq_length, 4]					,					self.range_bbox					)
        # Ensure that bbox is legal
        for i in range(bbox.shape[0]					):
              for j in range(bbox.shape[1]					):
                    if bbox[i, j, 3] < bbox[i, j, 1]:
                          _lowerCAmelCase		   =   bbox[i, j, 3]
                          _lowerCAmelCase		   =   bbox[i, j, 1]
                          _lowerCAmelCase		   =   t
                    if bbox[i, j, 2] < bbox[i, j, 0]:
                          _lowerCAmelCase		   =   bbox[i, j, 2]
                          _lowerCAmelCase		   =   bbox[i, j, 0]
                          _lowerCAmelCase		   =   t

        _lowerCAmelCase		   =   None
        if self.use_input_mask:
              _lowerCAmelCase		   =   ids_tensor([self.batch_size, self.seq_length]					,					vocab_size=2					)

        _lowerCAmelCase		   =   None
        if self.use_token_type_ids:
              _lowerCAmelCase		   =   ids_tensor([self.batch_size, self.seq_length]					,					self.type_vocab_size					)

        _lowerCAmelCase		   =   None
        _lowerCAmelCase		   =   None
        _lowerCAmelCase		   =   None
        if self.use_labels:
              _lowerCAmelCase		   =   ids_tensor([self.batch_size]					,					self.type_sequence_label_size					)
              _lowerCAmelCase		   =   ids_tensor([self.batch_size, self.seq_length]					,					self.num_labels					)
              _lowerCAmelCase		   =   ids_tensor([self.batch_size]					,					self.num_choices					)

        _lowerCAmelCase		   =   self.get_config()

        return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



  def   a   (  self   :					str					):


        """simple docstring"""



        return LayoutLMConfig(
            vocab_size=self.vocab_size					,					hidden_size=self.hidden_size					,					num_hidden_layers=self.num_hidden_layers					,					num_attention_heads=self.num_attention_heads					,					intermediate_size=self.intermediate_size					,					hidden_act=self.hidden_act					,					hidden_dropout_prob=self.hidden_dropout_prob					,					attention_probs_dropout_prob=self.attention_probs_dropout_prob					,					max_position_embeddings=self.max_position_embeddings					,					type_vocab_size=self.type_vocab_size					,					initializer_range=self.initializer_range					,					)



  def   a   (  self   :					List[Any]					,					__lowerCAmelCase   :					Optional[Any]					,					__lowerCAmelCase   :					List[str]					,					__lowerCAmelCase   :					Tuple					,					__lowerCAmelCase   :					Union[str, Any]					,					__lowerCAmelCase   :					Optional[Any]					,					__lowerCAmelCase   :					Union[str, Any]					,					__lowerCAmelCase   :					str					,					__lowerCAmelCase   :					Any					):


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMModel(config=lowerCamelCase__					)
        model.to(lowerCamelCase__					)
        model.eval()
        _lowerCAmelCase		   =   model(lowerCamelCase__					,					lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					)
        _lowerCAmelCase		   =   model(lowerCamelCase__					,					lowerCamelCase__					,					token_type_ids=lowerCamelCase__					)
        _lowerCAmelCase		   =   model(lowerCamelCase__					,					lowerCamelCase__					)
        self.parent.assertEqual(result.last_hidden_state.shape					,					(self.batch_size, self.seq_length, self.hidden_size)					)
        self.parent.assertEqual(result.pooler_output.shape					,					(self.batch_size, self.hidden_size)					)



  def   a   (  self   :					Any					,					__lowerCAmelCase   :					Dict					,					__lowerCAmelCase   :					Tuple					,					__lowerCAmelCase   :					str					,					__lowerCAmelCase   :					str					,					__lowerCAmelCase   :					List[str]					,					__lowerCAmelCase   :					str					,					__lowerCAmelCase   :					Tuple					,					__lowerCAmelCase   :					List[Any]					):


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMForMaskedLM(config=lowerCamelCase__					)
        model.to(lowerCamelCase__					)
        model.eval()
        _lowerCAmelCase		   =   model(lowerCamelCase__					,					lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					,					labels=lowerCamelCase__					)
        self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.seq_length, self.vocab_size)					)



  def   a   (  self   :					Any					,					__lowerCAmelCase   :					str					,					__lowerCAmelCase   :					List[Any]					,					__lowerCAmelCase   :					str					,					__lowerCAmelCase   :					Optional[int]					,					__lowerCAmelCase   :					List[Any]					,					__lowerCAmelCase   :					Any					,					__lowerCAmelCase   :					Optional[Any]					,					__lowerCAmelCase   :					Optional[Any]					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.num_labels
        _lowerCAmelCase		   =   LayoutLMForSequenceClassification(lowerCamelCase__					)
        model.to(lowerCamelCase__					)
        model.eval()
        _lowerCAmelCase		   =   model(
            lowerCamelCase__					,					lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					,					labels=lowerCamelCase__					)
        self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.num_labels)					)



  def   a   (  self   :					List[str]					,					__lowerCAmelCase   :					int					,					__lowerCAmelCase   :					List[str]					,					__lowerCAmelCase   :					Optional[Any]					,					__lowerCAmelCase   :					Union[str, Any]					,					__lowerCAmelCase   :					List[str]					,					__lowerCAmelCase   :					List[Any]					,					__lowerCAmelCase   :					str					,					__lowerCAmelCase   :					Any					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.num_labels
        _lowerCAmelCase		   =   LayoutLMForTokenClassification(config=lowerCamelCase__					)
        model.to(lowerCamelCase__					)
        model.eval()
        _lowerCAmelCase		   =   model(lowerCamelCase__					,					lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					,					labels=lowerCamelCase__					)
        self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.seq_length, self.num_labels)					)



  def   a   (  self   :					int					,					__lowerCAmelCase   :					List[str]					,					__lowerCAmelCase   :					Dict					,					__lowerCAmelCase   :					Tuple					,					__lowerCAmelCase   :					Tuple					,					__lowerCAmelCase   :					List[Any]					,					__lowerCAmelCase   :					Optional[int]					,					__lowerCAmelCase   :					Tuple					,					__lowerCAmelCase   :					int					):


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMForQuestionAnswering(config=lowerCamelCase__					)
        model.to(lowerCamelCase__					)
        model.eval()
        _lowerCAmelCase		   =   model(
            lowerCamelCase__					,					bbox=lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					,					start_positions=lowerCamelCase__					,					end_positions=lowerCamelCase__					,					)
        self.parent.assertEqual(result.start_logits.shape					,					(self.batch_size, self.seq_length)					)
        self.parent.assertEqual(result.end_logits.shape					,					(self.batch_size, self.seq_length)					)



  def   a   (  self   :					Dict					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.prepare_config_and_inputs()
        (
            (
            _lowerCAmelCase						
        )   ,  (
            _lowerCAmelCase						
        )   ,  (
            _lowerCAmelCase						
        )   ,  (
            _lowerCAmelCase						
        )   ,  (
            _lowerCAmelCase						
        )   ,  (
            _lowerCAmelCase						
        )   ,  (
            _lowerCAmelCase						
        )   ,  (
            _lowerCAmelCase						
        )   ,  						
        )		   =   config_and_inputs
        _lowerCAmelCase		   =   {
            'input_ids': input_ids,
            'bbox': bbox,
            'token_type_ids': token_type_ids,
            'attention_mask': input_mask,
        }
        return config, inputs_dict





@require_torch
class        SCREAMING_SNAKE_CASE   (				__UpperCAmelCase	,						__UpperCAmelCase	,						unittest.TestCase ):
  """simple docstring"""







  __A     =		(
      (
          LayoutLMModel,
          LayoutLMForMaskedLM,
          LayoutLMForSequenceClassification,
          LayoutLMForTokenClassification,
          LayoutLMForQuestionAnswering,
      )
      if is_torch_available()
      else None
  )
  __A     =		(
      {
          """document-question-answering""": LayoutLMForQuestionAnswering,
          """feature-extraction""": LayoutLMModel,
          """fill-mask""": LayoutLMForMaskedLM,
          """text-classification""": LayoutLMForSequenceClassification,
          """token-classification""": LayoutLMForTokenClassification,
          """zero-shot""": LayoutLMForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  __A     =		True



  def   a   (  self   :					List[str]					):


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMModelTester(self					)
        _lowerCAmelCase		   =   ConfigTester(self					,					config_class=lowerCamelCase__					,					hidden_size=37					)



  def   a   (  self   :					Dict					):


        """simple docstring"""



        self.config_tester.run_common_tests()



  def   a   (  self   :					Optional[int]					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCamelCase__					)



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


        """simple docstring"""



        _lowerCAmelCase		   =   self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
              _lowerCAmelCase		   =   type
              self.model_tester.create_and_check_model(*lowerCamelCase__					)



  def   a   (  self   :					str					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase__					)



  def   a   (  self   :					List[str]					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase__					)



  def   a   (  self   :					Optional[int]					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__					)



  def   a   (  self   :					Dict					):


        """simple docstring"""



        _lowerCAmelCase		   =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*lowerCamelCase__					)



  @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'					)
  def   a   (  self   :					Tuple					):


        """simple docstring"""



        pass





def 			A_   ( ):
      # Here we prepare a batch of 2 sequences to test a LayoutLM forward pass on:
      # fmt: off
      _lowerCAmelCase		   =   torch.tensor([[101,1_019,1_014,1_016,1_037,12_849,4_747,1_004,14_246,2_278,5_439,4_524,5_002,2_930,2_193,2_930,4_341,3_208,1_005,1_055,2_171,2_848,11_300,3_531,102],[101,4_070,4_034,7_020,1_024,3_058,1_015,1_013,2_861,1_013,6_070,19_274,2_772,6_205,27_814,16_147,16_147,4_343,2_047,10_283,10_969,14_389,1_012,2_338,102]]  ,  device=lowercase__    )  # noqa: E231
      _lowerCAmelCase		   =   torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]  ,  device=lowercase__    )  # noqa: E231
      _lowerCAmelCase		   =   torch.tensor([[[0,0,0,0],[423,237,440,251],[427,272,441,287],[419,115,437,129],[961,885,992,912],[256,38,330,58],[256,38,330,58],[336,42,353,57],[360,39,401,56],[360,39,401,56],[411,39,471,59],[479,41,528,59],[533,39,630,60],[67,113,134,131],[141,115,209,132],[68,149,133,166],[141,149,187,164],[195,148,287,165],[195,148,287,165],[195,148,287,165],[295,148,349,165],[441,149,492,166],[497,149,546,164],[64,201,125,218],[1_000,1_000,1_000,1_000]],[[0,0,0,0],[662,150,754,166],[665,199,742,211],[519,213,554,228],[519,213,554,228],[134,433,187,454],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[314,469,376,482],[504,684,582,706],[941,825,973,900],[941,825,973,900],[941,825,973,900],[941,825,973,900],[610,749,652,765],[130,659,168,672],[176,657,237,672],[238,657,312,672],[443,653,628,672],[443,653,628,672],[716,301,825,317],[1_000,1_000,1_000,1_000]]]  ,  device=lowercase__    )  # noqa: E231
      _lowerCAmelCase		   =   torch.tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]  ,  device=lowercase__    )  # noqa: E231
      # these are sequence labels (i.e. at the token level)
      _lowerCAmelCase		   =   torch.tensor([[-100,10,10,10,9,1,-100,7,7,-100,7,7,4,2,5,2,8,8,-100,-100,5,0,3,2,-100],[-100,12,12,12,-100,12,10,-100,-100,-100,-100,10,12,9,-100,-100,-100,10,10,10,9,12,-100,10,-100]]  ,  device=lowercase__    )  # noqa: E231
      # fmt: on

      return input_ids, attention_mask, bbox, token_type_ids, labels





@require_torch
class        SCREAMING_SNAKE_CASE   (				unittest.TestCase ):
  """simple docstring"""







  @slow
  def   a   (  self   :					int					):


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMModel.from_pretrained('microsoft/layoutlm-base-uncased'					).to(lowerCamelCase__					)

        _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase		   =   prepare_layoutlm_batch_inputs()

        # forward pass
        _lowerCAmelCase		   =   model(input_ids=lowerCamelCase__					,					bbox=lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					)

        # test the sequence output on [0, :3, :3]
        _lowerCAmelCase		   =   torch.tensor(
            [[0.1_7_8_5, -0.1_9_4_7, -0.0_4_2_5], [-0.3_2_5_4, -0.2_8_0_7, 0.2_5_5_3], [-0.5_3_9_1, -0.3_3_2_2, 0.3_3_6_4]]					,					device=lowerCamelCase__					,					)

        self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]					,					lowerCamelCase__					,					atol=1e-3					)					)

        # test the pooled output on [1, :3]
        _lowerCAmelCase		   =   torch.tensor([-0.6_5_8_0, -0.0_2_1_4, 0.8_5_5_2]					,					device=lowerCamelCase__					)

        self.assertTrue(torch.allclose(outputs.pooler_output[1, :3]					,					lowerCamelCase__					,					atol=1e-3					)					)



  @slow
  def   a   (  self   :					Tuple					):


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMForSequenceClassification.from_pretrained('microsoft/layoutlm-base-uncased'					,					num_labels=2					).to(
            lowerCamelCase__					)

        _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase		   =   prepare_layoutlm_batch_inputs()

        # forward pass
        _lowerCAmelCase		   =   model(
            input_ids=lowerCamelCase__					,					bbox=lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					,					labels=torch.tensor([1, 1]					,					device=lowerCamelCase__					)					,					)

        # test whether we get a loss as a scalar
        _lowerCAmelCase		   =   outputs.loss
        _lowerCAmelCase		   =   torch.Size([]					)
        self.assertEqual(loss.shape					,					lowerCamelCase__					)

        # test the shape of the logits
        _lowerCAmelCase		   =   outputs.logits
        _lowerCAmelCase		   =   torch.Size((2, 2)					)
        self.assertEqual(logits.shape					,					lowerCamelCase__					)



  @slow
  def   a   (  self   :					Optional[Any]					):


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased'					,					num_labels=13					).to(
            lowerCamelCase__					)

        _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase		   =   prepare_layoutlm_batch_inputs()

        # forward pass
        _lowerCAmelCase		   =   model(
            input_ids=lowerCamelCase__					,					bbox=lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					,					labels=lowerCamelCase__					)

        # test the loss calculation to be around 2.65
        # expected_loss = torch.tensor(2.65, device=torch_device)

        # The loss is currently somewhat random and can vary between 0.1-0.3 atol.
        # self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=0.1))

        # test the shape of the logits
        _lowerCAmelCase		   =   outputs.logits
        _lowerCAmelCase		   =   torch.Size((2, 25, 13)					)
        self.assertEqual(logits.shape					,					lowerCamelCase__					)



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


        """simple docstring"""



        _lowerCAmelCase		   =   LayoutLMForQuestionAnswering.from_pretrained('microsoft/layoutlm-base-uncased'					).to(lowerCamelCase__					)

        _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase   ,  _lowerCAmelCase		   =   prepare_layoutlm_batch_inputs()

        # forward pass
        _lowerCAmelCase		   =   model(input_ids=lowerCamelCase__					,					bbox=lowerCamelCase__					,					attention_mask=lowerCamelCase__					,					token_type_ids=lowerCamelCase__					)

        # test the shape of the logits
        _lowerCAmelCase		   =   torch.Size((2, 25)					)
        self.assertEqual(outputs.start_logits.shape					,					lowerCamelCase__					)
        self.assertEqual(outputs.end_logits.shape					,					lowerCamelCase__					)






'''simple docstring'''




import argparse

import runhouse as rh
import torch
from nlp_example import training_function

from accelerate.utils import PrepareForLaunch, patch_environment







def 			A_   ( *_lowerCamelCase			:							Union[str, Any]    ):
      _lowerCAmelCase		   =   torch.cuda.device_count()
      print(F"Device count: {num_processes}"    )
      with patch_environment(
          world_size=_lowerCamelCase  ,  master_addr='127.0.01'  ,  master_port='29500'  ,  mixed_precision=args[1].mixed_precision    ):
            _lowerCAmelCase		   =   PrepareForLaunch(_lowerCamelCase  ,  distributed_type='MULTI_GPU'    )
            torch.multiprocessing.start_processes(_lowerCamelCase  ,  args=_lowerCamelCase  ,  nprocs=_lowerCamelCase  ,  start_method='spawn'    )


if __name__ == "__main__":
 # Refer to https://runhouse-docs.readthedocs-hosted.com/en/main/rh_primitives/cluster.html#hardware-setup
 # for cloud access setup instructions (if using on-demand hardware), and for API specifications.

 # on-demand GPU
 # gpu = rh.cluster(name='rh-cluster', instance_type='V100:1', provider='cheapest', use_spot=False)  # single GPU
 snake_case  =			rh.cluster(name='''rh-cluster''', instance_type='''V100:4''', provider='''cheapest''', use_spot=False)  # multi GPU
 gpu.up_if_not()

 # on-prem GPU
 # gpu = rh.cluster(
 #           ips=["ip_addr"], ssh_creds={ssh_user:"<username>", ssh_private_key:"<key_path>"}, name="rh-cluster"
 #       )

 # Set up remote function
 snake_case  =			[
     '''pip:./''',
     '''transformers''',
     '''datasets''',
     '''evaluate''',
     '''tqdm''',
     '''scipy''',
     '''scikit-learn''',
     '''tensorboard''',
     '''torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117''',
 ]
 snake_case  =			rh.function(fn=launch_train, system=gpu, reqs=reqs, name='''train_bert_glue''')

 # Define train args/config, run train function
 snake_case  =			argparse.Namespace(cpu=False, mixed_precision='''fp16''')
 snake_case  =			{'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
 launch_train_gpu(config, train_args, stream_logs=True)

 # Alternatively, we can just run as instructed in the README (but only because there's already a wrapper CLI):
 # gpu.install_packages(reqs)
 # gpu.run(['accelerate launch --multi_gpu accelerate/examples/nlp_example.py'])





