




'''simple docstring'''





import unittest

import numpy as np

from transformers import (
    GPTSanJapaneseConfig,
    GPTSanJapaneseForConditionalGeneration,
    GPTSanJapaneseModel,
    GPTSanJapaneseTokenizer,
    is_torch_available,
)
from transformers.generation import GenerationConfig
from transformers.testing_utils import require_torch, slow, tooslow, 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




class      UpperCamelCase       :



  def __init__(			self		,     __lowerCAmelCase		,     __lowerCAmelCase=3_6000		,     __lowerCAmelCase=13		,     __lowerCAmelCase=7		,     __lowerCAmelCase=True		,     __lowerCAmelCase=32		,     __lowerCAmelCase=42		,     __lowerCAmelCase=5		,     __lowerCAmelCase=2		,     __lowerCAmelCase=4		,     __lowerCAmelCase=2		,     __lowerCAmelCase=32		,     __lowerCAmelCase=80		,     __lowerCAmelCase=33		,     __lowerCAmelCase=0.0		,     __lowerCAmelCase=1e-6		,     __lowerCAmelCase=100		,     __lowerCAmelCase=0.0		,     ):







      '''simple docstring'''


      _a				:							Tuple							       =      vocab_size
      _a				:							Dict							       =      parent
      _a				:							int							       =      batch_size
      _a				:							List[str]							       =      num_contexts
      # For common tests
      _a				:							Optional[Any]							       =      self.num_contexts
      _a				:							Optional[int]							       =      is_training
      _a				:							Dict							       =      hidden_size
      _a				:							Optional[int]							       =      num_ext_layers
      _a				:							Optional[int]							       =      ext_size
      _a				:							Tuple							       =      num_hidden_layers
      _a				:							int							       =      num_attention_heads
      _a				:							Optional[int]							       =      num_experts
      _a				:							Union[str, Any]							       =      d_ff
      _a				:							int							       =      d_ext
      _a				:							int							       =      d_spout
      _a				:							int							       =      dropout_rate
      _a				:							int							       =      layer_norm_epsilon
      _a				:							Dict							       =      expert_capacity
      _a				:							str							       =      router_jitter_noise



  def     a_							(			self ):







      '''simple docstring'''


      return GPTSanJapaneseConfig.from_pretrained("""Tanrei/GPTSAN-japanese""" )



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Dict							       =      ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size )

      _a				:							Any							       =      self.get_config()

      return (config, input_ids)



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Union[str, Any]							       =      ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size )

      _a				:							Tuple							       =      self.get_config()

      return (config, {"input_ids": input_ids})



  def     a_							(			self ):







      '''simple docstring'''


      return GPTSanJapaneseConfig(
          vocab_size=self.vocab_size		,     num_contexts=self.seq_length		,     d_model=self.hidden_size		,     d_ff=self.d_ff		,     d_ext=self.d_ext		,     d_spout=self.d_spout		,     num_switch_layers=self.num_hidden_layers - self.num_ext_layers		,     num_ext_layers=self.num_ext_layers		,     num_heads=self.num_attention_heads		,     num_experts=self.num_experts		,     expert_capacity=self.expert_capacity		,     dropout_rate=self.dropout_rate		,     layer_norm_epsilon=self.layer_norm_epsilon		,     router_jitter_noise=self.router_jitter_noise		,     )



  def     a_							(			self		,     __lowerCAmelCase		,     __lowerCAmelCase		,     ):







      '''simple docstring'''


      _a				:							Dict							       =      GPTSanJapaneseForConditionalGeneration(config=__lowerCAmelCase )
      model.to(__lowerCAmelCase )
      model.eval()
      _a				:							Dict							       =      model(
          input_ids=__lowerCAmelCase		,     )
      self.parent.assertIsNotNone(__lowerCAmelCase )






@require_torch
class      UpperCamelCase       (	snake_case__      ,		snake_case__      ,		unittest.TestCase       ):
  a     :Optional[Any]          =					(GPTSanJapaneseModel,) if is_torch_available() else ()
  a     :Union[str, Any]          =					(
      {
          'conversational': GPTSanJapaneseForConditionalGeneration,
          'feature-extraction': GPTSanJapaneseForConditionalGeneration,
          'summarization': GPTSanJapaneseForConditionalGeneration,
          'text2text-generation': GPTSanJapaneseForConditionalGeneration,
          'translation': GPTSanJapaneseForConditionalGeneration,
      }
      if is_torch_available()
      else {}
  )
  a     :Optional[int]          =					False
  a     :int          =					False
  a     :Union[str, Any]          =					False
  a     :Tuple          =					False
  a     :str          =					False
  a     :Optional[int]          =					False
  a     :Union[str, Any]          =					False
  a     :Optional[Any]          =					False
  # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
  a     :str          =					[0.8, 0.9]



  def     a_							(			self		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


      if pipeline_test_casse_name == "SummarizationPipelineTests":
          # TODO: fix `_reorder_cache` is not implemented for this model
          return True
      elif pipeline_test_casse_name == "Text2TextGenerationPipelineTests":
          # TODO: check this.
          return True

      return False



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Union[str, Any]							       =      GPTSanJapaneseTester(self )
      _a				:							Optional[int]							       =      ConfigTester(self		,     config_class=__lowerCAmelCase		,     d_model=37 )



  def     a_							(			self ):







      '''simple docstring'''


      GPTSanJapaneseConfig()



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Optional[int]							       =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__lowerCAmelCase )



  @unittest.skip(
      reason="""skip for now as the computed `max_memory` by `model_split_percents` in the test method will be changed inside `from_pretrained`""" )
  def     a_							(			self ):







      '''simple docstring'''


      super().test_model_parallelism()



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







      '''simple docstring'''


      pass






@require_torch
class      UpperCamelCase       (	snake_case__      ,		snake_case__      ,		unittest.TestCase       ):
  a     :int          =					(GPTSanJapaneseForConditionalGeneration,) if is_torch_available() else ()
  a     :Any          =					False
  a     :Tuple          =					False
  a     :Optional[int]          =					False
  a     :Dict          =					False
  a     :List[str]          =					False
  a     :str          =					False
  # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
  a     :Optional[int]          =					[0.8, 0.9]



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Any							       =      GPTSanJapaneseTester(self )
      _a				:							Union[str, Any]							       =      ConfigTester(self		,     config_class=__lowerCAmelCase		,     d_model=37 )



  def     a_							(			self ):







      '''simple docstring'''


      GPTSanJapaneseConfig()



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Union[str, Any]							       =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__lowerCAmelCase )



  @unittest.skip(
      reason="""skip for now as the computed `max_memory` by `model_split_percents` in the test method will be changed inside `from_pretrained`""" )
  def     a_							(			self ):







      '''simple docstring'''


      super().test_model_parallelism()



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







      '''simple docstring'''


      pass



  @slow
  def     a_							(			self ):







      '''simple docstring'''


      _a				:							List[str]							       =      GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							List[Any]							       =      GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							Union[str, Any]							       =      tokenizer.encode("""武田信玄は"""		,     return_tensors="""pt""" )
      _a				:							List[str]							       =      model(__lowerCAmelCase )
      _a				:							int							       =      outputs.logits.detach().cpu().numpy()
      # Output of original model created with mesh-tensoflow
      _a				:							Union[str, Any]							       =      [
          # fmt: off
          [-12.037839889526367, -12.433061599731445, -14.333840370178223, -12.450345993041992, -11.1661376953125,
          -11.930137634277344, -10.659740447998047, -12.909574508666992, -13.241043090820312, -13.398579597473145,
          -11.107524871826172, -12.3685941696167, -22.97943115234375, -10.481067657470703, -12.484030723571777,
          -12.807360649108887, -14.769700050354004, -12.233579635620117, -13.428145408630371, -22.624177932739258],
          [-7.511149883270264, -8.281851768493652, -7.943127155303955, -7.55021333694458, -6.49869966506958,
          -7.586796283721924, -6.978085994720459, -7.839145183563232, -8.21964168548584, -8.695091247558594,
          -6.706910610198975, -6.6585798263549805, -19.565698623657227, -5.353842735290527, -8.350686073303223,
          -8.039388656616211, -10.856569290161133, -7.75154447555542, -8.819022178649902, -19.51532745361328],
          [-9.73066234588623, -10.223922729492188, -9.932981491088867, -11.857836723327637, -7.662626266479492,
          -11.13529109954834, -7.765097618103027, -11.472923278808594, -9.543149948120117, -11.905633926391602,
          -9.366164207458496, -11.5734281539917, -23.699003219604492, -9.429590225219727, -10.42839241027832,
          -10.585240364074707, -10.94771957397461, -11.095416069030762, -10.390240669250488, -23.769372940063477],
          [-9.728265762329102, -9.859712600708008, -10.09729290008545, -9.678522109985352, -6.879519939422607,
          -9.68487548828125, -4.2803425788879395, -10.018914222717285, -9.308445930480957, -10.63394546508789,
          -8.083646774291992, -9.06301498413086, -21.904266357421875, -8.90160846710205, -8.841876029968262,
          -11.856719970703125, -12.079398155212402, -11.233753204345703, -10.177338600158691, -21.87256622314453],
          [-9.669764518737793, -9.614198684692383, -9.814510345458984, -9.996501922607422, -11.375690460205078,
          -10.113405227661133, -10.546867370605469, -10.04369068145752, -10.907809257507324, -10.504216194152832,
          -11.129199028015137, -10.151124000549316, -21.96586799621582, -9.086349487304688, -11.730339050292969,
          -10.460667610168457, -10.298049926757812, -10.784148216247559, -10.840693473815918, -22.03152847290039],
          # fmt: on
      ]
      _a				:							int							       =      np.array(__lowerCAmelCase ).flatten()
      _a				:							List[str]							       =      output_logits[0, :, :20].flatten()

      def check(__lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase=5e-4 ):
          return abs(a - b ) < epsilon * max(abs(__lowerCAmelCase )		,     abs(__lowerCAmelCase ) )

      self.assertTrue(np.all([check(target[i]		,     predict[i] ) for i in range(len(__lowerCAmelCase ) )] ) )



  @slow
  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Any							       =      GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							int							       =      GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      model.to(__lowerCAmelCase )

      # set deterministically
      _a				:							Union[str, Any]							       =      GenerationConfig.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							Optional[int]							       =      1

      # use different length sentences to test batching
      _a				:							Optional[Any]							       =      [
          """甲斐なら武田と言うほど""",
          """織田信長は、""",
      ]

      _a				:							Union[str, Any]							       =      """left"""
      _a				:							List[Any]							       =      tokenizer(__lowerCAmelCase		,     return_tensors="""pt"""		,     padding=__lowerCAmelCase )
      _a				:							Union[str, Any]							       =      inputs["""input_ids"""].to(__lowerCAmelCase )

      self.assertNotEqual(inputs["""attention_mask"""][0].numpy().tolist()		,     inputs["""attention_mask"""][1].numpy().tolist() )

      _a				:							Any							       =      model.generate(
          input_ids=__lowerCAmelCase		,     attention_mask=inputs["""attention_mask"""].to(__lowerCAmelCase )		,     max_new_tokens=3		,     generation_config=__lowerCAmelCase		,     )

      _a				:							List[str]							       =      tokenizer(sentences[0]		,     return_tensors="""pt""" ).input_ids.to(__lowerCAmelCase )
      _a				:							Dict							       =      model.generate(
          input_ids=__lowerCAmelCase		,     max_new_tokens=3		,     generation_config=__lowerCAmelCase )

      _a				:							Optional[int]							       =      tokenizer(sentences[1]		,     return_tensors="""pt""" ).input_ids.to(__lowerCAmelCase )
      _a				:							List[str]							       =      model.generate(input_ids=__lowerCAmelCase		,     max_new_tokens=3		,     generation_config=__lowerCAmelCase )

      self.assertNotEqual(inputs_non_padded.shape		,     inputs_padded.shape )

      _a				:							int							       =      tokenizer.batch_decode(__lowerCAmelCase		,     skip_special_tokens=__lowerCAmelCase )
      _a				:							str							       =      tokenizer.decode(output_non_padded[0]		,     skip_special_tokens=__lowerCAmelCase )
      _a				:							int							       =      tokenizer.decode(output_padded[0]		,     skip_special_tokens=__lowerCAmelCase )

      _a				:							Dict							       =      [
          """甲斐なら武田と言うほど甲斐の武田""",
          """織田信長は、このような""",
      ]
      self.assertListEqual(__lowerCAmelCase		,     __lowerCAmelCase )
      self.assertListEqual(__lowerCAmelCase		,     [non_padded_sentence, padded_sentence] )



  @tooslow
  def     a_							(			self ):







      '''simple docstring'''


      _a				:							int							       =      GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							Tuple							       =      GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      # Output of original model created with mesh-tensoflow
      _a				:							Any							       =      [
          ("""武田信玄は""", 3_5675),
          ("""武田信玄は、""", 45),
          ("""武田信玄は、この""", 29),
          ("""武田信玄は、このよう""", 3_0642),
          ("""武田信玄は、このような""", 3_5680),
          ("""武田信玄は、このような「""", 8640),
          ("""武田信玄は、このような「武田""", 3_1617),
          ("""武田信玄は、このような「武田家""", 3_0646),
          ("""武田信玄は、このような「武田家の""", 3_1617),
          ("""武田信玄は、このような「武田家の家""", 3_1381),
      ]
      for input, output in target:
          _a				:							List[Any]							       =      tokenizer.encode(__lowerCAmelCase		,     return_tensors="""pt""" )
          _a				:							Any							       =      model(__lowerCAmelCase )
          _a				:							Dict							       =      outputs.logits.detach().cpu().numpy()[0]
          _a				:							Tuple							       =      np.argmax(output_logits[-1] )
          self.assertEqual(__lowerCAmelCase		,     __lowerCAmelCase )



  @slow
  def     a_							(			self ):







      '''simple docstring'''


      _a				:							int							       =      GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							int							       =      GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      model.to(__lowerCAmelCase )

      # set deterministically
      _a				:							Optional[int]							       =      GenerationConfig.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							str							       =      1

      _a				:							Any							       =      """武田信玄は、"""
      _a				:							List[str]							       =      tokenizer(__lowerCAmelCase		,     return_tensors="""pt""" ).input_ids.to(__lowerCAmelCase )
      _a				:							Dict							       =      tokenizer([input_text, input_text]		,     return_tensors="""pt""" ).input_ids.to(__lowerCAmelCase )

      # spout from uniform and one-hot
      _a				:							int							       =      [
          # fmt: off
          [0.87882208, 0.38426396, 0.33220248, 0.43890406, 0.16562252,
          0.04803985, 0.211572  , 0.23188473, 0.37153068, 0.7836377 ,
          0.02160172, 0.38761719, 0.75290772, 0.90198857, 0.34365777,
          0.64168169, 0.44318471, 0.14575746, 0.92562881, 0.40812148,
          0.29019122, 0.88861599, 0.65524846, 0.43563456, 0.38177187,
          0.70832965, 0.81527892, 0.68832812, 0.38833192, 0.4561522 ,
          0.14828817, 0.47248213, 0.54357335, 0.82009566, 0.1338884 ,
          0.02755417, 0.19764677, 0.2422084 , 0.04757674, 0.65409606,
          0.0824589 , 0.03304383, 0.94387689, 0.98764509, 0.82433901,
          0.27646741, 0.64907493, 0.76009406, 0.30087915, 0.17904689,
          0.41601714, 0.67046398, 0.10422822, 0.08447374, 0.07354344,
          0.61423565, 0.70284866, 0.7532333 , 0.1972038 , 0.29575659,
          0.90583886, 0.29265307, 0.50000175, 0.70407655, 0.889363  ,
          0.81904418, 0.66829128, 0.64468815, 0.56563723, 0.85601875,
          0.94924672, 0.00166762, 0.25220643, 0.74540219, 0.67993247,
          0.1549675 , 0.39385352, 0.92153607, 0.63745931, 0.27759043,
          0.84702295, 0.65904271, 0.58676614, 0.8666936 , 0.39607438,
          0.79954983, 0.42220697, 0.39650381, 0.7849864 , 0.56150201,
          0.15678925, 0.14746032, 0.34542114, 0.47026783, 0.11956489,
          0.25421435, 0.33788901, 0.68934842, 0.36424685, 0.71737898,
          0.38983449, 0.94393779, 0.39575588, 0.36616553, 0.87104665,
          0.64630203, 0.22516905, 0.88270804, 0.15031338, 0.75144345,
          0.46459025, 0.85396454, 0.86355643, 0.65139851, 0.70266061,
          0.30241389, 0.81056497, 0.88865969, 0.38773807, 0.70635849,
          0.90718459, 0.43245789, 0.28000654, 0.45935562, 0.08773519,
          0.9552151 , 0.93901511, 0.22489288], # uniform
          [1., 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., 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., 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.],
          # fmt: on
      ]

      _a				:							Union[str, Any]							       =      model.generate(
          input_ids=__lowerCAmelCase		,     spout=spouts[0]		,     max_new_tokens=20		,     generation_config=__lowerCAmelCase		,     )

      _a				:							Optional[Any]							       =      model.generate(
          input_ids=__lowerCAmelCase		,     spout=spouts[1]		,     max_new_tokens=20		,     generation_config=__lowerCAmelCase		,     )

      _a				:							List[Any]							       =      model.generate(
          input_ids=__lowerCAmelCase		,     spout=__lowerCAmelCase		,     max_new_tokens=20		,     generation_config=__lowerCAmelCase		,     )

      _a				:							int							       =      tokenizer.decode(outputa[0] )
      _a				:							Dict							       =      tokenizer.decode(outputa[0] )
      _a				:							List[str]							       =      tokenizer.batch_decode(__lowerCAmelCase )

      _a				:							Any							       =      [
          """武田信玄は、武田氏の滅亡後、武田氏の居城であった甲斐武田氏の居城である""",
          """武田信玄は、武田家の滅亡を防ぐため、武田家の家臣である武田信虎を討""",
      ]
      self.assertListEqual(__lowerCAmelCase		,     __lowerCAmelCase )
      self.assertListEqual(__lowerCAmelCase		,     [outa_sentence, outa_sentence] )



  @slow
  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Optional[int]							       =      GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							Union[str, Any]							       =      GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      model.to(__lowerCAmelCase )

      # set deterministically
      _a				:							List[str]							       =      GenerationConfig.from_pretrained("""Tanrei/GPTSAN-japanese""" )
      _a				:							Optional[Any]							       =      1

      _a				:							Dict							       =      """武田信玄"""
      _a				:							Optional[Any]							       =      """織田信長"""
      _a				:							List[Any]							       =      """は、"""
      _a				:							Union[str, Any]							       =      """が、"""
      _a				:							str							       =      tokenizer(__lowerCAmelCase		,     prefix_text=__lowerCAmelCase		,     return_tensors="""pt""" )
      _a				:							List[Any]							       =      tokenizer(__lowerCAmelCase		,     prefix_text=__lowerCAmelCase		,     return_tensors="""pt""" )
      _a				:							str							       =      tokenizer([[prefix_text_a, input_text_a], [prefix_text_a, input_text_a]]		,     return_tensors="""pt""" )

      _a				:							List[Any]							       =      model.generate(
          input_ids=input_tok_a.input_ids.to(__lowerCAmelCase )		,     token_type_ids=input_tok_a.token_type_ids.to(__lowerCAmelCase )		,     max_new_tokens=20		,     generation_config=__lowerCAmelCase		,     )

      _a				:							int							       =      model.generate(
          input_ids=input_tok_a.input_ids.to(__lowerCAmelCase )		,     token_type_ids=input_tok_a.token_type_ids.to(__lowerCAmelCase )		,     max_new_tokens=20		,     generation_config=__lowerCAmelCase		,     )

      _a				:							Optional[int]							       =      model.generate(
          input_ids=input_tok_a.input_ids.to(__lowerCAmelCase )		,     token_type_ids=input_tok_a.token_type_ids.to(__lowerCAmelCase )		,     attention_mask=input_tok_a.attention_mask.to(__lowerCAmelCase )		,     max_new_tokens=20		,     generation_config=__lowerCAmelCase		,     )

      _a				:							Tuple							       =      tokenizer.decode(outputa[0] )
      _a				:							Optional[int]							       =      tokenizer.decode(outputa[0] )
      _a				:							List[str]							       =      tokenizer.batch_decode(__lowerCAmelCase )

      _a				:							List[Any]							       =      [
          """武田信玄は、武田氏の祖である武田信虎を、その子・武田信友を擁して""",
          """織田信長が、織田信長の妻・お市の方を妻として迎えたという逸話が残""",
      ]
      self.assertListEqual(__lowerCAmelCase		,     __lowerCAmelCase )
      self.assertListEqual(__lowerCAmelCase		,     [outa_sentence, outa_sentence] )








'''simple docstring'''


import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
     import torch

     from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
     import pyanvml.pyanvml as nvml


UpperCamelCase			     =					logging.get_logger(__name__)






class 						lowerCAmelCase_	(       UpperCAmelCase_							):







  '''simple docstring'''

  UpperCamelCase_			:      PyTorchBenchmarkArguments
  UpperCamelCase_			:      PretrainedConfig
  UpperCamelCase_			:      str										=				"PyTorch"



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







       '''simple docstring'''



       return torch.__version__



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		float:







       '''simple docstring'''



       A:		List[str]          =		self._prepare_inference_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_speed(_inference	)



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		[Memory, Optional[MemorySummary]]:







       '''simple docstring'''



       A:		Any          =		self._prepare_inference_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_memory(_inference	)



  def 					_snake_case     ( self				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		float:







       '''simple docstring'''



       A:		List[Any]          =		self._prepare_train_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_speed(_train	)



  def 					_snake_case     ( self				:							int					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		[Memory, Optional[MemorySummary]]:







       '''simple docstring'''



       A:		List[Any]          =		self._prepare_train_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_memory(_train	)



  def 					_snake_case     ( self				:							Union[str, Any]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		Callable[[], None]:







       '''simple docstring'''



       A:		int          =		self.config_dict[model_name]

       if self.args.torchscript:
            A:		int          =		True

       A:		Any          =		(
           hasattr(SCREAMING_SNAKE_CASE_					,  '''architectures'''	)
           and isinstance(config.architectures					,  SCREAMING_SNAKE_CASE_	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 A:		Optional[int]          =		config.architectures[0]
                 A:		Optional[int]          =		__import__('''transformers'''					,  fromlist=[model_class]	)
                 A:		Dict          =		getattr(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
                 A:		Optional[Any]          =		model_cls(SCREAMING_SNAKE_CASE_	)
            except ImportError:
                 raise ImportError(
                     f"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                     ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
            A:		Union[str, Any]          =		MODEL_MAPPING[config.__class__](SCREAMING_SNAKE_CASE_	)

       model.eval()
       model.to(self.args.device	)

       # encoder-decoder has vocab size saved differently
       A:		Tuple          =		config.vocab_size if hasattr(SCREAMING_SNAKE_CASE_					,  '''vocab_size'''	) else config.encoder.vocab_size
       A:		List[Any]          =		torch.randint(SCREAMING_SNAKE_CASE_					,  (batch_size, sequence_length)					,  dtype=torch.long					,  device=self.args.device	)

       if self.args.fpaa:
            logger.info('''Running training in Mixed Precision...'''	)
            if not self.args.is_gpu:
                 raise ValueError('''Mixed precision is possible only for GPU.'''	)
            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       if self.args.torchscript:
            with torch.no_grad():
                 A:		Any          =		torch.jit.trace(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       else:
            A:		Dict          =		model

       def encoder_decoder_forward():
            with torch.no_grad():
                 A:		List[str]          =		inference_model(SCREAMING_SNAKE_CASE_					,  decoder_input_ids=SCREAMING_SNAKE_CASE_	)
            return outputs

       def encoder_forward():
            with torch.no_grad():
                 A:		Union[str, Any]          =		inference_model(SCREAMING_SNAKE_CASE_	)
            return outputs

       A:		List[Any]          =		encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
       return _forward



  def 					_snake_case     ( self				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		Callable[[], None]:







       '''simple docstring'''



       A:		int          =		self.config_dict[model_name]

       A:		str          =		(
           hasattr(SCREAMING_SNAKE_CASE_					,  '''architectures'''	)
           and isinstance(config.architectures					,  SCREAMING_SNAKE_CASE_	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 A:		List[str]          =		config.architectures[0]
                 A:		List[str]          =		__import__('''transformers'''					,  fromlist=[model_class]	)
                 A:		Optional[int]          =		getattr(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
                 A:		Union[str, Any]          =		model_cls(SCREAMING_SNAKE_CASE_	)
            except ImportError:
                 raise ImportError(
                     f"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                     ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
            A:		Union[str, Any]          =		MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE_	)

       if self.args.torchscript:
            raise NotImplementedError('''Training for torchscript is currently not implemented'''	)
       else:
            A:		str          =		model

       model.train()
       model.to(self.args.device	)

       # encoder-decoder has vocab size saved differently
       A:		Tuple          =		config.vocab_size if hasattr(SCREAMING_SNAKE_CASE_					,  '''vocab_size'''	) else config.encoder.vocab_size
       A:		str          =		torch.randint(SCREAMING_SNAKE_CASE_					,  (batch_size, sequence_length)					,  dtype=torch.long					,  device=self.args.device	)

       if self.args.fpaa:
            logger.info('''Running training in Mixed Precision...'''	)
            if not self.args.is_gpu:
                 raise ValueError('''Mixed precision is possible only for GPU.'''	)

            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       def compute_loss_and_backprob_encoder():
            A:		str          =		train_model(SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)[0]
            loss.backward()
            return loss

       def compute_loss_and_backprob_encoder_decoder():
            A:		Dict          =		train_model(SCREAMING_SNAKE_CASE_					,  decoder_input_ids=SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)[0]
            loss.backward()
            return loss

       A:		str          =		(
           compute_loss_and_backprob_encoder_decoder
           if config.is_encoder_decoder
           else compute_loss_and_backprob_encoder
       )
       return _train



  def 					_snake_case     ( self				:							Any					,  SCREAMING_SNAKE_CASE_				:							Any	)			->		float:







       '''simple docstring'''



       try:
            if self.args.is_tpu or self.args.torchscript:
                 # run additional 10 times to stabilize compilation for tpu and torchscript
                 logger.info('''Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'''	)
                 timeit.repeat(
                     SCREAMING_SNAKE_CASE_					,  repeat=1					,  number=5					,  )

            # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
            A:		Any          =		timeit.repeat(
                SCREAMING_SNAKE_CASE_					,  repeat=self.args.repeat					,  number=10					,  )

            if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                 import torch_xla.debug.metrics as met

                 self.print_fn(met.metrics_report()	)

            return min(SCREAMING_SNAKE_CASE_	) / 10.0
       except RuntimeError as e:
            self.print_fn(f"""Doesn't fit on GPU. {e}"""	)
            return "N/A"




  def 					_snake_case     ( self				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							Callable[[], None]	)			->		[Memory, MemorySummary]:







       '''simple docstring'''



       try:
            if self.args.trace_memory_line_by_line:
                 A:		List[Any]          =		start_memory_tracing('''transformers'''	)

            if self.args.is_tpu:
                 # tpu
                 raise NotImplementedError(
                     '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'''
                     ''' `--no-memory` or `args.memory=False`'''	)
            elif self.args.is_gpu:
                 if not is_pyanvml_available():
                      logger.warning(
                          '''py3nvml not installed, we won\'t log GPU memory usage. '''
                          '''Install py3nvml (pip install py3nvml) to log information about GPU.'''	)
                      A:		Optional[int]          =		'''N/A'''
                 else:
                      logger.info(
                          '''Measuring total GPU usage on GPU device. Make sure to not have additional processes running'''
                          ''' on the same GPU.'''	)
                      # init nvml
                      nvml.nvmlInit()
                      func()
                      A:		Union[str, Any]          =		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
                      A:		Union[str, Any]          =		nvml.nvmlDeviceGetMemoryInfo(SCREAMING_SNAKE_CASE_	)
                      A:		Tuple          =		meminfo.used
                      A:		Tuple          =		Memory(SCREAMING_SNAKE_CASE_	)
                      # shutdown nvml
                      nvml.nvmlShutdown()
            else:
                 # cpu
                 A:		int          =		measure_peak_memory_cpu(SCREAMING_SNAKE_CASE_	)
                 A:		str          =		Memory(SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	) else memory_bytes

            if self.args.trace_memory_line_by_line:
                 A:		Dict          =		stop_memory_tracing(SCREAMING_SNAKE_CASE_	)
            else:
                 A:		List[str]          =		None

            return memory, summary
       except RuntimeError as e:
            self.print_fn(f"""Doesn't fit on GPU. {e}"""	)
            return "N/A", None






