

import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

__UpperCAmelCase						=					[
    '''cross_validation.py''',
    '''gradient_accumulation.py''',
    '''local_sgd.py''',
    '''multi_process_metrics.py''',
    '''memory.py''',
    '''automatic_gradient_accumulation.py''',
    '''fsdp_with_peak_mem_tracking.py''',
    '''deepspeed_with_config_support.py''',
    '''megatron_lm_gpt_pretraining.py''',
]



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






  def 			SCREAMING_SNAKE_CASE       (     self					:       str			,				lowerCAmelCase					:       Union[str, Any]			,				lowerCAmelCase					:       int			,				lowerCAmelCase					:       Optional[int] = None			,				lowerCAmelCase					:       Optional[Any] = None       )				->      Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase    :     Optional[Any]          =		None
     __lowerCAmelCase    :     int          =		os.path.abspath(os.path.join("""examples"""			,				"""by_feature"""       )       )
     __lowerCAmelCase    :     str          =		os.path.abspath("""examples"""       )
     for item in os.listdir(_snake_case       ):
        if item not in EXCLUDE_EXAMPLES:
           __lowerCAmelCase    :     Dict          =		os.path.join(_snake_case			,				_snake_case       )
           if os.path.isfile(_snake_case       ) and ".py" in item_path:
              with self.subTest(
                  tested_script=_snake_case			,				feature_script=_snake_case			,				tested_section="""main()""" if parser_only else """training_function()"""			,				):
                 __lowerCAmelCase    :     List[Any]          =		compare_against_test(
                     os.path.join(_snake_case			,				_snake_case       )			,				_snake_case			,				_snake_case			,				_snake_case       )
                 __lowerCAmelCase    :     List[str]          =		"""\n""".join(_snake_case       )
                 if special_strings is not None:
                    for string in special_strings:
                       __lowerCAmelCase    :     List[Any]          =		diff.replace(_snake_case			,				""""""       )
                 self.assertEqual(_snake_case			,				""""""       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       Dict       )				->      List[str]:






     """simple docstring"""







     self.one_complete_example("""complete_nlp_example.py"""			,				_snake_case       )
     self.one_complete_example("""complete_nlp_example.py"""			,				_snake_case       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       Union[str, Any]       )				->      Tuple:






     """simple docstring"""







     __lowerCAmelCase    :     Tuple          =		os.path.abspath(os.path.join("""examples"""			,				"""cv_example.py"""       )       )
     __lowerCAmelCase    :     Union[str, Any]          =		[
         """ """ * 16 + """{\n\n""",
         """ """ * 20 + """\"accuracy\": eval_metric[\"accuracy\"],\n\n""",
         """ """ * 20 + """\"f1\": eval_metric[\"f1\"],\n\n""",
         """ """ * 20 + """\"train_loss\": total_loss.item() / len(train_dataloader),\n\n""",
         """ """ * 20 + """\"epoch\": epoch,\n\n""",
         """ """ * 16 + """},\n\n""",
         """ """ * 16 + """step=epoch,\n""",
         """ """ * 12,
         """ """ * 8 + """for step, batch in enumerate(active_dataloader):\n""",
     ]
     self.one_complete_example("""complete_cv_example.py"""			,				_snake_case			,				_snake_case			,				_snake_case       )
     self.one_complete_example("""complete_cv_example.py"""			,				_snake_case			,				_snake_case			,				_snake_case       )






@mock.patch.dict(os.environ  ,				{"TESTING_MOCKED_DATALOADERS": "1"}					)
class SCREAMING_SNAKE_CASE      (    a__					):
  """simple docstring"""






  lowerCamelCase       :     Optional[int]		 =False







  @classmethod
  def 			SCREAMING_SNAKE_CASE       (     cls					:       Union[str, Any]       )				->      Optional[int]:






     """simple docstring"""







     super().setUpClass()
     __lowerCAmelCase    :     int          =		tempfile.mkdtemp()
     __lowerCAmelCase    :     str          =		os.path.join(cls._tmpdir			,				"""default_config.yml"""       )

     write_basic_config(save_location=cls.configPath       )
     __lowerCAmelCase    :     str          =		["""accelerate""", """launch""", """--config_file""", cls.configPath]







  @classmethod
  def 			SCREAMING_SNAKE_CASE       (     cls					:       Dict       )				->      Union[str, Any]:






     """simple docstring"""







     super().tearDownClass()
     shutil.rmtree(cls._tmpdir       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       List[str]       )				->      Tuple:






     """simple docstring"""







     __lowerCAmelCase    :     List[Any]          =		f'''\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps epoch\n        --output_dir {self.tmpdir}\n        '''.split()
     run_command(self._launch_args + testargs       )
     self.assertTrue(os.path.exists(os.path.join(self.tmpdir			,				"""epoch_0"""       )       )       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       Any       )				->      Tuple:






     """simple docstring"""







     __lowerCAmelCase    :     Optional[int]          =		f'''\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps 1\n        --output_dir {self.tmpdir}\n        '''.split()
     __lowerCAmelCase    :     Any          =		run_command(self._launch_args + testargs       )
     self.assertTrue(os.path.exists(os.path.join(self.tmpdir			,				"""step_2"""       )       )       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       Any       )				->      Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase    :     Optional[int]          =		f'''\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir			,				'epoch_0'       )}\n        '''.split()
     __lowerCAmelCase    :     Optional[int]          =		run_command(self._launch_args + testargs			,				return_stdout=_snake_case       )
     self.assertNotIn("""epoch 0:"""			,				_snake_case       )
     self.assertIn("""epoch 1:"""			,				_snake_case       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       Dict       )				->      Tuple:






     """simple docstring"""







     __lowerCAmelCase    :     Optional[Any]          =		f'''\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir			,				'step_2'       )}\n        '''.split()
     __lowerCAmelCase    :     int          =		run_command(self._launch_args + testargs			,				return_stdout=_snake_case       )
     if torch.cuda.is_available():
        __lowerCAmelCase    :     Union[str, Any]          =		torch.cuda.device_count()
     else:
        __lowerCAmelCase    :     int          =		1
     if num_processes > 1:
        self.assertNotIn("""epoch 0:"""			,				_snake_case       )
        self.assertIn("""epoch 1:"""			,				_snake_case       )
     else:
        self.assertIn("""epoch 0:"""			,				_snake_case       )
        self.assertIn("""epoch 1:"""			,				_snake_case       )







  @slow
  def 			SCREAMING_SNAKE_CASE       (     self					:       int       )				->      List[str]:






     """simple docstring"""







     __lowerCAmelCase    :     List[Any]          =		"""\n        examples/by_feature/cross_validation.py\n        --num_folds 2\n        """.split()
     with mock.patch.dict(os.environ			,				{"""TESTING_MOCKED_DATALOADERS""": """0"""}       ):
        __lowerCAmelCase    :     int          =		run_command(self._launch_args + testargs			,				return_stdout=_snake_case       )
        __lowerCAmelCase    :     Tuple          =		re.findall("""({.+})"""			,				_snake_case       )
        __lowerCAmelCase    :     Optional[int]          =		[r for r in results if """accuracy""" in r][-1]
        __lowerCAmelCase    :     Optional[Any]          =		ast.literal_eval(_snake_case       )
        self.assertGreaterEqual(results["""accuracy"""]			,				0.75       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       List[Any]       )				->      int:






     """simple docstring"""







     __lowerCAmelCase    :     List[Any]          =		["""examples/by_feature/multi_process_metrics.py"""]
     run_command(self._launch_args + testargs       )







  @require_trackers
  @mock.patch.dict(os.environ			,				{"""WANDB_MODE""": """offline"""}       )
  def 			SCREAMING_SNAKE_CASE       (     self					:       Dict       )				->      int:






     """simple docstring"""







     with tempfile.TemporaryDirectory() as tmpdir:
        __lowerCAmelCase    :     Dict          =		f'''\n            examples/by_feature/tracking.py\n            --with_tracking\n            --project_dir {tmpdir}\n            '''.split()
        run_command(self._launch_args + testargs       )
        self.assertTrue(os.path.exists(os.path.join(_snake_case			,				"""tracking"""       )       )       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       Optional[Any]       )				->      str:






     """simple docstring"""







     __lowerCAmelCase    :     List[str]          =		["""examples/by_feature/gradient_accumulation.py"""]
     run_command(self._launch_args + testargs       )







  def 			SCREAMING_SNAKE_CASE       (     self					:       str       )				->      List[Any]:






     """simple docstring"""







     __lowerCAmelCase    :     Optional[int]          =		["""examples/by_feature/local_sgd.py"""]
     run_command(self._launch_args + testargs       )









"""simple docstring"""
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, 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 import (
           OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
           OpenAIGPTConfig,
           OpenAIGPTDoubleHeadsModel,
           OpenAIGPTForSequenceClassification,
           OpenAIGPTLMHeadModel,
           OpenAIGPTModel,
       )







class       a							:



 def __init__(     self				,       _snake_case				,       _snake_case=13				,       _snake_case=7				,       _snake_case=True				,       _snake_case=True				,       _snake_case=True				,       _snake_case=99				,       _snake_case=32				,       _snake_case=5				,       _snake_case=4				,       _snake_case=37				,       _snake_case="gelu"				,       _snake_case=0.1				,       _snake_case=0.1				,       _snake_case=5_12				,       _snake_case=16				,       _snake_case=2				,       _snake_case=0.02				,       _snake_case=3				,       _snake_case=4				,       _snake_case=None				,       ):







  """simple docstring"""






  lowerCAmelCase    				= parent
  lowerCAmelCase    				= batch_size
  lowerCAmelCase    				= seq_length
  lowerCAmelCase    				= is_training
  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    				= self.vocab_size - 1



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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

  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    				= OpenAIGPTConfig(
      vocab_size=self.vocab_size				,       n_embd=self.hidden_size				,       n_layer=self.num_hidden_layers				,       n_head=self.num_attention_heads				,       n_positions=self.max_position_embeddings				,       pad_token_id=self.pad_token_id				,       )

  lowerCAmelCase    				= ids_tensor([self.num_hidden_layers, self.num_attention_heads]				,       2     )

  return (
      config,
      input_ids,
      head_mask,
      token_type_ids,
      sequence_labels,
      token_labels,
      choice_labels,
  )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTModel(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       head_mask=_snake_case     )
  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case     )
  lowerCAmelCase    				= model(_snake_case     )

  self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTLMHeadModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.loss.shape				,       ()     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTDoubleHeadsModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.loss.shape				,       ()     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.num_labels
  lowerCAmelCase    				= OpenAIGPTForSequenceClassification(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= ids_tensor([self.batch_size]				,       self.type_sequence_label_size     )
  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.prepare_config_and_inputs()
  (
      (
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,
  )    				= config_and_inputs
  lowerCAmelCase    				= {
      'input_ids': input_ids,
      'token_type_ids': token_type_ids,
      'head_mask': head_mask,
  }

  return config, inputs_dict







@require_torch
class       a							(   a__   , a__   , a__   , unittest.TestCase  ):
 snake_case__             =       (
     (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
     if is_torch_available()
     else ()
 )
 snake_case__             =       (
     (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
 )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
 snake_case__             =       (
     {
         '''feature-extraction''': OpenAIGPTModel,
         '''text-classification''': OpenAIGPTForSequenceClassification,
         '''text-generation''': OpenAIGPTLMHeadModel,
         '''zero-shot''': OpenAIGPTForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
   # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
   # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
   # tiny config could not be created.
   return True

  return False



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case=False     ):







  """simple docstring"""






  lowerCAmelCase    				= super()._prepare_for_class(_snake_case				,       _snake_case				,       return_labels=_snake_case     )

  if return_labels:
   if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
    lowerCAmelCase    				= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)				,       dtype=torch.long				,       device=_snake_case				,       )
    lowerCAmelCase    				= inputs_dict['labels']
    lowerCAmelCase    				= inputs_dict['labels']
    lowerCAmelCase    				= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices)				,       dtype=torch.long				,       device=_snake_case				,       )
    lowerCAmelCase    				= torch.zeros(
        self.model_tester.batch_size				,       dtype=torch.long				,       device=_snake_case     )
  return inputs_dict



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTModelTester(self     )
  lowerCAmelCase    				= ConfigTester(self				,       config_class=_snake_case				,       n_embd=37     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  self.config_tester.run_common_tests()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_lm_head_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_double_lm_head_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*_snake_case     )



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   lowerCAmelCase    				= OpenAIGPTModel.from_pretrained(_snake_case     )
   self.assertIsNotNone(_snake_case     )







@require_torch
class       a							(   unittest.TestCase  ):



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTLMHeadModel.from_pretrained('openai-gpt'     )
  model.to(_snake_case     )
  lowerCAmelCase    				= torch.tensor([[4_81, 47_35, 5_44]]				,       dtype=torch.long				,       device=_snake_case     )  # the president is
  lowerCAmelCase    				= [
      4_81,
      47_35,
      5_44,
      2_46,
      9_63,
      8_70,
      7_62,
      2_39,
      2_44,
      4_04_77,
      2_44,
      2_49,
      7_19,
      8_81,
      4_87,
      5_44,
      2_40,
      2_44,
      6_03,
      4_81,
  ]  # the president is a very good man. " \n " i\'m sure he is, " said the

  lowerCAmelCase    				= model.generate(_snake_case				,       do_sample=_snake_case     )
  self.assertListEqual(output_ids[0].tolist()				,       _snake_case     )





