





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 	_lowercase		:

 def __init__(				self			:							List[str]     ,   snake_case			:							List[Any]     ,   snake_case			:							Tuple=3_6_0_0_0     ,   snake_case			:							Any=1_3     ,   snake_case			:							str=7     ,   snake_case			:							Optional[int]=True     ,   snake_case			:							str=3_2     ,   snake_case			:							Dict=4_2     ,   snake_case			:							Tuple=5     ,   snake_case			:							str=2     ,   snake_case			:							Optional[int]=4     ,   snake_case			:							List[str]=2     ,   snake_case			:							str=3_2     ,   snake_case			:							str=8_0     ,   snake_case			:							Any=3_3     ,   snake_case			:							Dict=0.0     ,   snake_case			:							Optional[int]=1e-6     ,   snake_case			:							List[Any]=1_0_0     ,   snake_case			:							List[str]=0.0     ,   )     ->		Any:

  """simple docstring"""





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

 def 		SCREAMING_SNAKE_CASE__     (				self			:							Any    )     ->		int:

  """simple docstring"""





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

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

  """simple docstring"""





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

  UpperCamelCase_      :		Union[str, Any]				=      self.get_config()

  return (config, input_ids)

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

  """simple docstring"""





  UpperCamelCase_      :		Any				=      ids_tensor([self.batch_size, self.seq_length]     ,   self.vocab_size    )

  UpperCamelCase_      :		int				=      self.get_config()

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

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

  """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 		SCREAMING_SNAKE_CASE__     (				self			:							str     ,   snake_case			:							List[Any]     ,   snake_case			:							Tuple     ,   )     ->		str:

  """simple docstring"""





  UpperCamelCase_      :		int				=      GPTSanJapaneseForConditionalGeneration(config=lowercase_    )
  model.to(lowercase_    )
  model.eval()
  UpperCamelCase_      :		str				=      model(
      input_ids=lowercase_     ,   )
  self.parent.assertIsNotNone(lowercase_    )





@require_torch
class 	_lowercase		(      __UpperCAmelCase		,   __UpperCAmelCase		,   unittest.TestCase ):
 lowercase         =     (GPTSanJapaneseModel,) if is_torch_available() else ()
 lowercase         =     (
     {
         'conversational': GPTSanJapaneseForConditionalGeneration,
         'feature-extraction': GPTSanJapaneseForConditionalGeneration,
         'summarization': GPTSanJapaneseForConditionalGeneration,
         'text2text-generation': GPTSanJapaneseForConditionalGeneration,
         'translation': GPTSanJapaneseForConditionalGeneration,
     }
     if is_torch_available()
     else {}
 )
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
 lowercase         =     [0.8, 0.9]

 def 		SCREAMING_SNAKE_CASE__     (				self			:							str     ,   snake_case			:							str     ,   snake_case			:							List[str]     ,   snake_case			:							List[Any]     ,   snake_case			:							int     ,   snake_case			:							Dict    )     ->		List[Any]:

  """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 		SCREAMING_SNAKE_CASE__     (				self			:							Optional[int]    )     ->		List[str]:

  """simple docstring"""





  UpperCamelCase_      :		List[str]				=      GPTSanJapaneseTester(self    )
  UpperCamelCase_      :		str				=      ConfigTester(self     ,   config_class=lowercase_     ,   d_model=3_7    )

 def 		SCREAMING_SNAKE_CASE__     (				self			:							int    )     ->		int:

  """simple docstring"""





  GPTSanJapaneseConfig()

 def 		SCREAMING_SNAKE_CASE__     (				self			:							Tuple    )     ->		Dict:

  """simple docstring"""





  UpperCamelCase_      :		int				=      self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*lowercase_    )

 @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 		SCREAMING_SNAKE_CASE__     (				self			:							List[Any]    )     ->		Optional[int]:

  """simple docstring"""





  super().test_model_parallelism()


 @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'    )
 def 		SCREAMING_SNAKE_CASE__     (				self			:							str    )     ->		Optional[int]:

  """simple docstring"""





  pass





@require_torch
class 	_lowercase		(      __UpperCAmelCase		,   __UpperCAmelCase		,   unittest.TestCase ):
 lowercase         =     (GPTSanJapaneseForConditionalGeneration,) if is_torch_available() else ()
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 lowercase         =     False
 # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
 lowercase         =     [0.8, 0.9]

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

  """simple docstring"""





  UpperCamelCase_      :		List[str]				=      GPTSanJapaneseTester(self    )
  UpperCamelCase_      :		str				=      ConfigTester(self     ,   config_class=lowercase_     ,   d_model=3_7    )

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

  """simple docstring"""





  GPTSanJapaneseConfig()

 def 		SCREAMING_SNAKE_CASE__     (				self			:							Any    )     ->		List[Any]:

  """simple docstring"""





  UpperCamelCase_      :		Optional[Any]				=      self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*lowercase_    )

 @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 		SCREAMING_SNAKE_CASE__     (				self			:							Optional[Any]    )     ->		Optional[int]:

  """simple docstring"""





  super().test_model_parallelism()

 @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'    )
 def 		SCREAMING_SNAKE_CASE__     (				self			:							Dict    )     ->		List[Any]:

  """simple docstring"""





  pass

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

  """simple docstring"""





  UpperCamelCase_      :		Union[str, Any]				=      GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		List[str]				=      GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		Optional[int]				=      tokenizer.encode('武田信玄は'     ,   return_tensors='pt'    )
  UpperCamelCase_      :		List[str]				=      model(lowercase_    )
  UpperCamelCase_      :		List[Any]				=      outputs.logits.detach().cpu().numpy()
  # Output of original model created with mesh-tensoflow
  UpperCamelCase_      :		List[str]				=      [
      # 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
  ]
  UpperCamelCase_      :		List[str]				=      np.array(lowercase_    ).flatten()
  UpperCamelCase_      :		List[Any]				=      output_logits[0, :, :2_0].flatten()

  def check(snake_case			:							Dict     ,   snake_case			:							Any     ,   snake_case			:							str=5e-4    ):
   return abs(a - b    ) < epsilon * max(abs(lowercase_    )     ,   abs(lowercase_    )    )

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

 @slow
 def 		SCREAMING_SNAKE_CASE__     (				self			:							Any    )     ->		Optional[int]:

  """simple docstring"""





  UpperCamelCase_      :		Any				=      GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		List[Any]				=      GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'    )
  model.to(lowercase_    )

  # set deterministically
  UpperCamelCase_      :		List[Any]				=      GenerationConfig.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		List[Any]				=      1

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

  UpperCamelCase_      :		str				=      '''left'''
  UpperCamelCase_      :		Optional[int]				=      tokenizer(lowercase_     ,   return_tensors='pt'     ,   padding=lowercase_    )
  UpperCamelCase_      :		Optional[int]				=      inputs['''input_ids'''].to(lowercase_    )

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

  UpperCamelCase_      :		Optional[Any]				=      model.generate(
      input_ids=lowercase_     ,   attention_mask=inputs['attention_mask'].to(lowercase_    )     ,   max_new_tokens=3     ,   generation_config=lowercase_     ,   )

  UpperCamelCase_      :		Dict				=      tokenizer(sentences[0]     ,   return_tensors='pt'    ).input_ids.to(lowercase_    )
  UpperCamelCase_      :		List[str]				=      model.generate(
      input_ids=lowercase_     ,   max_new_tokens=3     ,   generation_config=lowercase_    )

  UpperCamelCase_      :		Optional[int]				=      tokenizer(sentences[1]     ,   return_tensors='pt'    ).input_ids.to(lowercase_    )
  UpperCamelCase_      :		int				=      model.generate(input_ids=lowercase_     ,   max_new_tokens=3     ,   generation_config=lowercase_    )

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

  UpperCamelCase_      :		List[str]				=      tokenizer.batch_decode(lowercase_     ,   skip_special_tokens=lowercase_    )
  UpperCamelCase_      :		Any				=      tokenizer.decode(output_non_padded[0]     ,   skip_special_tokens=lowercase_    )
  UpperCamelCase_      :		Union[str, Any]				=      tokenizer.decode(output_padded[0]     ,   skip_special_tokens=lowercase_    )

  UpperCamelCase_      :		Any				=      [
      '''甲斐なら武田と言うほど甲斐の武田''',
      '''織田信長は、このような''',
  ]
  self.assertListEqual(lowercase_     ,   lowercase_    )
  self.assertListEqual(lowercase_     ,   [non_padded_sentence, padded_sentence]    )

 @tooslow
 def 		SCREAMING_SNAKE_CASE__     (				self			:							Optional[int]    )     ->		Any:

  """simple docstring"""





  UpperCamelCase_      :		Any				=      GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		Optional[int]				=      GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'    )
  # Output of original model created with mesh-tensoflow
  UpperCamelCase_      :		Optional[Any]				=      [
      ('''武田信玄は''', 3_5_6_7_5),
      ('''武田信玄は、''', 4_5),
      ('''武田信玄は、この''', 2_9),
      ('''武田信玄は、このよう''', 3_0_6_4_2),
      ('''武田信玄は、このような''', 3_5_6_8_0),
      ('''武田信玄は、このような「''', 8_6_4_0),
      ('''武田信玄は、このような「武田''', 3_1_6_1_7),
      ('''武田信玄は、このような「武田家''', 3_0_6_4_6),
      ('''武田信玄は、このような「武田家の''', 3_1_6_1_7),
      ('''武田信玄は、このような「武田家の家''', 3_1_3_8_1),
  ]
  for input, output in target:
   UpperCamelCase_      :		Any				=      tokenizer.encode(lowercase_     ,   return_tensors='pt'    )
   UpperCamelCase_      :		Dict				=      model(lowercase_    )
   UpperCamelCase_      :		Dict				=      outputs.logits.detach().cpu().numpy()[0]
   UpperCamelCase_      :		int				=      np.argmax(output_logits[-1]    )
   self.assertEqual(lowercase_     ,   lowercase_    )

 @slow
 def 		SCREAMING_SNAKE_CASE__     (				self			:							Tuple    )     ->		Dict:

  """simple docstring"""





  UpperCamelCase_      :		List[str]				=      GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		Any				=      GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'    )
  model.to(lowercase_    )

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

  UpperCamelCase_      :		List[str]				=      '''武田信玄は、'''
  UpperCamelCase_      :		int				=      tokenizer(lowercase_     ,   return_tensors='pt'    ).input_ids.to(lowercase_    )
  UpperCamelCase_      :		Optional[Any]				=      tokenizer([input_text, input_text]     ,   return_tensors='pt'    ).input_ids.to(lowercase_    )

  # spout from uniform and one-hot
  UpperCamelCase_      :		List[str]				=      [
      # 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
  ]

  UpperCamelCase_      :		Dict				=      model.generate(
      input_ids=lowercase_     ,   spout=spouts[0]     ,   max_new_tokens=2_0     ,   generation_config=lowercase_     ,   )

  UpperCamelCase_      :		Tuple				=      model.generate(
      input_ids=lowercase_     ,   spout=spouts[1]     ,   max_new_tokens=2_0     ,   generation_config=lowercase_     ,   )

  UpperCamelCase_      :		str				=      model.generate(
      input_ids=lowercase_     ,   spout=lowercase_     ,   max_new_tokens=2_0     ,   generation_config=lowercase_     ,   )

  UpperCamelCase_      :		Dict				=      tokenizer.decode(outputa[0]    )
  UpperCamelCase_      :		Tuple				=      tokenizer.decode(outputa[0]    )
  UpperCamelCase_      :		Optional[Any]				=      tokenizer.batch_decode(lowercase_    )

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


 @slow
 def 		SCREAMING_SNAKE_CASE__     (				self			:							Tuple    )     ->		int:

  """simple docstring"""





  UpperCamelCase_      :		Any				=      GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		Optional[int]				=      GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'    )
  model.to(lowercase_    )

  # set deterministically
  UpperCamelCase_      :		List[Any]				=      GenerationConfig.from_pretrained('Tanrei/GPTSAN-japanese'    )
  UpperCamelCase_      :		Dict				=      1

  UpperCamelCase_      :		Any				=      '''武田信玄'''
  UpperCamelCase_      :		int				=      '''織田信長'''
  UpperCamelCase_      :		str				=      '''は、'''
  UpperCamelCase_      :		Dict				=      '''が、'''
  UpperCamelCase_      :		List[str]				=      tokenizer(lowercase_     ,   prefix_text=lowercase_     ,   return_tensors='pt'    )
  UpperCamelCase_      :		Dict				=      tokenizer(lowercase_     ,   prefix_text=lowercase_     ,   return_tensors='pt'    )
  UpperCamelCase_      :		Optional[Any]				=      tokenizer([[prefix_text_a, input_text_a], [prefix_text_a, input_text_a]]     ,   return_tensors='pt'    )

  UpperCamelCase_      :		Tuple				=      model.generate(
      input_ids=input_tok_a.input_ids.to(lowercase_    )     ,   token_type_ids=input_tok_a.token_type_ids.to(lowercase_    )     ,   max_new_tokens=2_0     ,   generation_config=lowercase_     ,   )

  UpperCamelCase_      :		List[Any]				=      model.generate(
      input_ids=input_tok_a.input_ids.to(lowercase_    )     ,   token_type_ids=input_tok_a.token_type_ids.to(lowercase_    )     ,   max_new_tokens=2_0     ,   generation_config=lowercase_     ,   )

  UpperCamelCase_      :		List[str]				=      model.generate(
      input_ids=input_tok_a.input_ids.to(lowercase_    )     ,   token_type_ids=input_tok_a.token_type_ids.to(lowercase_    )     ,   attention_mask=input_tok_a.attention_mask.to(lowercase_    )     ,   max_new_tokens=2_0     ,   generation_config=lowercase_     ,   )

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

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




from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


a_					      = logging.get_logger(__name__)  # pylint: disable=invalid-name

a_					      = '\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorEmb2EmbPipeline, KandinskyV22ControlnetImg2ImgPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)["depth"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline("depth-estimation")\n\n        >>> pipe_prior = KandinskyV22PriorEmb2EmbPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to("cuda")\n\n        >>> pipe = KandinskyV22ControlnetImg2ImgPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to("cuda")\n\n        >>> img = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/cat.png"\n        ... ).resize((768, 768))\n\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")\n\n        >>> prompt = "A robot, 4k photo"\n        >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"\n\n        >>> generator = torch.Generator(device="cuda").manual_seed(43)\n\n        >>> img_emb = pipe_prior(prompt=prompt, image=img, strength=0.85, generator=generator)\n        >>> negative_emb = pipe_prior(prompt=negative_prior_prompt, image=img, strength=1, generator=generator)\n\n        >>> images = pipe(\n        ...     image=img,\n        ...     strength=0.5,\n        ...     image_embeds=img_emb.image_embeds,\n        ...     negative_image_embeds=negative_emb.image_embeds,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save("robot_cat.png")\n        ```\n'





def      __lowercase  (		lowerCamelCase      :  str						,						lowerCamelCase      :  List[str]						,						lowerCamelCase      :  str=8			):
 UpperCamelCase_      :		Union[str, Any]				=      height // scale_factor**2
 if height % scale_factor**2 != 0:
  new_height += 1
 UpperCamelCase_      :		List[str]				=      width // scale_factor**2
 if width % scale_factor**2 != 0:
  new_width += 1
 return new_height * scale_factor, new_width * scale_factor





def      __lowercase  (		lowerCamelCase      :  Optional[Any]						,						lowerCamelCase      :  Optional[int]=512						,						lowerCamelCase      :  str=512			):
 UpperCamelCase_      :		Any				=      pil_image.resize((w, h)						,						resample=Image.BICUBIC						,						reducing_gap=1			)
 UpperCamelCase_      :		Tuple				=      np.array(pil_image.convert('RGB'			)			)
 UpperCamelCase_      :		Optional[int]				=      arr.astype(np.floataa			) / 1_2_7.5 - 1
 UpperCamelCase_      :		int				=      np.transpose(lowerCamelCase						,						[2, 0, 1]			)
 UpperCamelCase_      :		List[Any]				=      torch.from_numpy(lowerCamelCase			).unsqueeze(0			)
 return image






class 	_lowercase		(      snake_case_ ):

 def __init__(				self			:							Union[str, Any]     ,   snake_case			:							UNetaDConditionModel     ,   snake_case			:							DDPMScheduler     ,   snake_case			:							VQModel     ,   )     ->		str:

  """simple docstring"""





  super().__init__()

  self.register_modules(
      unet=snake_case     ,   scheduler=snake_case     ,   movq=snake_case     ,   )
  UpperCamelCase_      :		Optional[Any]				=      2 ** (len(self.movq.config.block_out_channels    ) - 1)

 def 		SCREAMING_SNAKE_CASE__     (				self			:							Dict     ,   snake_case			:							str     ,   snake_case			:							List[Any]     ,   snake_case			:							int    )     ->		str:

  """simple docstring"""





  UpperCamelCase_      :		Any				=      min(int(num_inference_steps * strength    )     ,   snake_case    )

  UpperCamelCase_      :		Optional[Any]				=      max(num_inference_steps - init_timestep     ,   0    )
  UpperCamelCase_      :		Union[str, Any]				=      self.scheduler.timesteps[t_start:]

  return timesteps, num_inference_steps - t_start

 def 		SCREAMING_SNAKE_CASE__     (				self			:							Dict     ,   snake_case			:							int     ,   snake_case			:							Union[str, Any]     ,   snake_case			:							List[str]     ,   snake_case			:							str     ,   snake_case			:							Optional[int]     ,   snake_case			:							Dict     ,   snake_case			:							List[Any]=None    )     ->		Union[str, Any]:

  """simple docstring"""





  if not isinstance(snake_case     ,   (torch.Tensor, PIL.Image.Image, list)    ):
   raise ValueError(
       f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(snake_case    )}"    )

  UpperCamelCase_      :		Union[str, Any]				=      image.to(device=snake_case     ,   dtype=snake_case    )

  UpperCamelCase_      :		Any				=      batch_size * num_images_per_prompt

  if image.shape[1] == 4:
   UpperCamelCase_      :		List[Any]				=      image

  else:
   if isinstance(snake_case     ,   snake_case    ) and len(snake_case    ) != batch_size:
    raise ValueError(
        f"You have passed a list of generators of length {len(snake_case    )}, but requested an effective batch"
        f" size of {batch_size}. Make sure the batch size matches the length of the generators."    )

   elif isinstance(snake_case     ,   snake_case    ):
    UpperCamelCase_      :		Any				=      [
        self.movq.encode(image[i : i + 1]    ).latent_dist.sample(generator[i]    ) for i in range(snake_case    )
    ]
    UpperCamelCase_      :		Union[str, Any]				=      torch.cat(snake_case     ,   dim=0    )
   else:
    UpperCamelCase_      :		Optional[Any]				=      self.movq.encode(snake_case    ).latent_dist.sample(snake_case    )

   UpperCamelCase_      :		List[Any]				=      self.movq.config.scaling_factor * init_latents

  UpperCamelCase_      :		List[str]				=      torch.cat([init_latents]     ,   dim=0    )

  UpperCamelCase_      :		Optional[Any]				=      init_latents.shape
  UpperCamelCase_      :		Optional[int]				=      randn_tensor(snake_case     ,   generator=snake_case     ,   device=snake_case     ,   dtype=snake_case    )

  # get latents
  UpperCamelCase_      :		Union[str, Any]				=      self.scheduler.add_noise(snake_case     ,   snake_case     ,   snake_case    )

  UpperCamelCase_      :		Optional[Any]				=      init_latents

  return latents

 def 		SCREAMING_SNAKE_CASE__     (				self			:							List[Any]     ,   snake_case			:							Dict=0    )     ->		Optional[Any]:

  """simple docstring"""





  if is_accelerate_available():
   from accelerate import cpu_offload
  else:
   raise ImportError('Please install accelerate via `pip install accelerate`'    )

  UpperCamelCase_      :		str				=      torch.device(f"cuda:{gpu_id}"    )

  UpperCamelCase_      :		Dict				=      [
      self.unet,
      self.movq,
  ]
  for cpu_offloaded_model in models:
   if cpu_offloaded_model is not None:
    cpu_offload(snake_case     ,   snake_case    )

 def 		SCREAMING_SNAKE_CASE__     (				self			:							str     ,   snake_case			:							Tuple=0    )     ->		Optional[Any]:

  """simple docstring"""





  if is_accelerate_available() and is_accelerate_version('>='     ,   '0.17.0.dev0'    ):
   from accelerate import cpu_offload_with_hook
  else:
   raise ImportError('`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'    )

  UpperCamelCase_      :		Any				=      torch.device(f"cuda:{gpu_id}"    )

  if self.device.type != "cpu":
   self.to('cpu'     ,   silence_dtype_warnings=snake_case    )
   torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

  UpperCamelCase_      :		Any				=      None
  for cpu_offloaded_model in [self.unet, self.movq]:
   UpperCamelCase_,				UpperCamelCase_      :		Any				=      cpu_offload_with_hook(snake_case     ,   snake_case     ,   prev_module_hook=snake_case    )

  # We'll offload the last model manually.
  UpperCamelCase_      :		Optional[int]				=      hook

 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def 		SCREAMING_SNAKE_CASE__     (				self			:							Union[str, Any]    )     ->		Optional[Any]:

  """simple docstring"""





  if not hasattr(self.unet     ,   '_hf_hook'    ):
   return self.device
  for module in self.unet.modules():
   if (
       hasattr(snake_case     ,   '_hf_hook'    )
       and hasattr(module._hf_hook     ,   'execution_device'    )
       and module._hf_hook.execution_device is not None
   ):
    return torch.device(module._hf_hook.execution_device    )
  return self.device


 @torch.no_grad()
 @replace_example_docstring(snake_case    )
 def __call__(				self			:							Optional[int]     ,   snake_case			:							Union[torch.FloatTensor, List[torch.FloatTensor]]     ,   snake_case			:							Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]]     ,   snake_case			:							Union[torch.FloatTensor, List[torch.FloatTensor]]     ,   snake_case			:							torch.FloatTensor     ,   snake_case			:							int = 5_1_2     ,   snake_case			:							int = 5_1_2     ,   snake_case			:							int = 1_0_0     ,   snake_case			:							float = 4.0     ,   snake_case			:							float = 0.3     ,   snake_case			:							int = 1     ,   snake_case			:							Optional[Union[torch.Generator, List[torch.Generator]]] = None     ,   snake_case			:							Optional[str] = "pil"     ,   snake_case			:							bool = True     ,   )     ->		Any:

  """simple docstring"""





  UpperCamelCase_      :		List[str]				=      self._execution_device

  UpperCamelCase_      :		str				=      guidance_scale > 1.0

  if isinstance(snake_case     ,   snake_case    ):
   UpperCamelCase_      :		Dict				=      torch.cat(snake_case     ,   dim=0    )
  if isinstance(snake_case     ,   snake_case    ):
   UpperCamelCase_      :		Optional[int]				=      torch.cat(snake_case     ,   dim=0    )
  if isinstance(snake_case     ,   snake_case    ):
   UpperCamelCase_      :		Dict				=      torch.cat(snake_case     ,   dim=0    )

  UpperCamelCase_      :		int				=      image_embeds.shape[0]

  if do_classifier_free_guidance:
   UpperCamelCase_      :		Optional[Any]				=      image_embeds.repeat_interleave(snake_case     ,   dim=0    )
   UpperCamelCase_      :		Any				=      negative_image_embeds.repeat_interleave(snake_case     ,   dim=0    )
   UpperCamelCase_      :		Tuple				=      hint.repeat_interleave(snake_case     ,   dim=0    )

  UpperCamelCase_      :		int				=      torch.cat([negative_image_embeds, image_embeds]     ,   dim=0    ).to(dtype=self.unet.dtype     ,   device=snake_case    )
  UpperCamelCase_      :		Union[str, Any]				=      torch.cat([hint, hint]     ,   dim=0    ).to(dtype=self.unet.dtype     ,   device=snake_case    )

  if not isinstance(snake_case     ,   snake_case    ):
   UpperCamelCase_      :		Optional[Any]				=      [image]
  if not all(isinstance(snake_case     ,   (PIL.Image.Image, torch.Tensor)    ) for i in image    ):
   raise ValueError(
       f"Input is in incorrect format: {[type(snake_case    ) for i in image]}. Currently, we only support  PIL image and pytorch tensor"    )

  UpperCamelCase_      :		Union[str, Any]				=      torch.cat([prepare_image(snake_case     ,   snake_case     ,   snake_case    ) for i in image]     ,   dim=0    )
  UpperCamelCase_      :		Dict				=      image.to(dtype=image_embeds.dtype     ,   device=snake_case    )

  UpperCamelCase_      :		List[Any]				=      self.movq.encode(snake_case    )['latents']
  UpperCamelCase_      :		str				=      latents.repeat_interleave(snake_case     ,   dim=0    )
  self.scheduler.set_timesteps(snake_case     ,   device=snake_case    )
  UpperCamelCase_,				UpperCamelCase_      :		int				=      self.get_timesteps(snake_case     ,   snake_case     ,   snake_case    )
  UpperCamelCase_      :		Any				=      timesteps[:1].repeat(batch_size * num_images_per_prompt    )
  UpperCamelCase_,				UpperCamelCase_      :		Optional[Any]				=      downscale_height_and_width(snake_case     ,   snake_case     ,   self.movq_scale_factor    )
  UpperCamelCase_      :		List[Any]				=      self.prepare_latents(
      snake_case     ,   snake_case     ,   snake_case     ,   snake_case     ,   image_embeds.dtype     ,   snake_case     ,   snake_case    )
  for i, t in enumerate(self.progress_bar(snake_case    )    ):
   # expand the latents if we are doing classifier free guidance
   UpperCamelCase_      :		int				=      torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents

   UpperCamelCase_      :		Optional[Any]				=      {'image_embeds': image_embeds, 'hint': hint}
   UpperCamelCase_      :		Optional[Any]				=      self.unet(
       sample=snake_case     ,   timestep=snake_case     ,   encoder_hidden_states=snake_case     ,   added_cond_kwargs=snake_case     ,   return_dict=snake_case     ,   )[0]

   if do_classifier_free_guidance:
    UpperCamelCase_,				UpperCamelCase_      :		Dict				=      noise_pred.split(latents.shape[1]     ,   dim=1    )
    UpperCamelCase_,				UpperCamelCase_      :		List[Any]				=      noise_pred.chunk(2    )
    UpperCamelCase_,				UpperCamelCase_      :		Union[str, Any]				=      variance_pred.chunk(2    )
    UpperCamelCase_      :		Union[str, Any]				=      noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    UpperCamelCase_      :		Tuple				=      torch.cat([noise_pred, variance_pred_text]     ,   dim=1    )

   if not (
       hasattr(self.scheduler.config     ,   'variance_type'    )
       and self.scheduler.config.variance_type in ["learned", "learned_range"]
   ):
    UpperCamelCase_,				UpperCamelCase_      :		int				=      noise_pred.split(latents.shape[1]     ,   dim=1    )

   # compute the previous noisy sample x_t -> x_t-1

   UpperCamelCase_      :		List[Any]				=      self.scheduler.step(
       snake_case     ,   snake_case     ,   snake_case     ,   generator=snake_case     ,   )[0]

  # post-processing
  UpperCamelCase_      :		Any				=      self.movq.decode(snake_case     ,   force_not_quantize=snake_case    )['sample']

  if output_type not in ["pt", "np", "pil"]:
   raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"    )

  if output_type in ["np", "pil"]:
   UpperCamelCase_      :		List[Any]				=      image * 0.5 + 0.5
   UpperCamelCase_      :		Optional[int]				=      image.clamp(0     ,   1    )
   UpperCamelCase_      :		Dict				=      image.cpu().permute(0     ,   2     ,   3     ,   1    ).float().numpy()

  if output_type == "pil":
   UpperCamelCase_      :		Optional[Any]				=      self.numpy_to_pil(snake_case    )

  if not return_dict:
   return (image,)

  return ImagePipelineOutput(images=snake_case    )



