



'''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  lowerCamelCase_   :
  """simple docstring"""




  def __init__(			self				:     Any      ,     _a				:     int      ,     _a				:     int=3_6000      ,     _a				:     Union[str, Any]=13      ,     _a				:     int=7      ,     _a				:     List[str]=True      ,     _a				:     int=32      ,     _a				:     Optional[int]=42      ,     _a				:     Optional[int]=5      ,     _a				:     List[str]=2      ,     _a				:     Optional[int]=4      ,     _a				:     Optional[int]=2      ,     _a				:     Optional[Any]=32      ,     _a				:     Any=80      ,     _a				:     Dict=33      ,     _a				:     Tuple=0.0      ,     _a				:     Union[str, Any]=1e-6      ,     _a				:     Any=100      ,     _a				:     List[Any]=0.0      ,     ) -> Optional[Any]:
    __lowerCamelCase	:	List[str]   	=			vocab_size
    __lowerCamelCase	:	str   	=			parent
    __lowerCamelCase	:	Optional[Any]   	=			batch_size
    __lowerCamelCase	:	Optional[Any]   	=			num_contexts
    # For common tests
    __lowerCamelCase	:	Optional[int]   	=			self.num_contexts
    __lowerCamelCase	:	List[Any]   	=			is_training
    __lowerCamelCase	:	Tuple   	=			hidden_size
    __lowerCamelCase	:	Optional[int]   	=			num_ext_layers
    __lowerCamelCase	:	Optional[int]   	=			ext_size
    __lowerCamelCase	:	str   	=			num_hidden_layers
    __lowerCamelCase	:	Any   	=			num_attention_heads
    __lowerCamelCase	:	str   	=			num_experts
    __lowerCamelCase	:	Optional[Any]   	=			d_ff
    __lowerCamelCase	:	Optional[Any]   	=			d_ext
    __lowerCamelCase	:	Union[str, Any]   	=			d_spout
    __lowerCamelCase	:	Optional[Any]   	=			dropout_rate
    __lowerCamelCase	:	Optional[Any]   	=			layer_norm_epsilon
    __lowerCamelCase	:	Any   	=			expert_capacity
    __lowerCamelCase	:	List[str]   	=			router_jitter_noise




  def 						_lowercase					(			self				:     Dict       ) -> Optional[int]:
    return GPTSanJapaneseConfig.from_pretrained('Tanrei/GPTSAN-japanese'       )




  def 						_lowercase					(			self				:     Optional[int]       ) -> Optional[int]:
    __lowerCamelCase	:	Any   	=			ids_tensor([self.batch_size, self.seq_length]      ,     self.vocab_size       )

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

    return (config, input_ids)




  def 						_lowercase					(			self				:     Union[str, Any]       ) -> str:
    __lowerCamelCase	:	Union[str, Any]   	=			ids_tensor([self.batch_size, self.seq_length]      ,     self.vocab_size       )

    __lowerCamelCase	:	List[str]   	=			self.get_config()

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




  def 						_lowercase					(			self				:     Tuple       ) -> List[Any]:
    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 						_lowercase					(			self				:     str      ,     _a				:     Any      ,     _a				:     Optional[int]      ,     ) -> List[str]:
    __lowerCamelCase	:	str   	=			GPTSanJapaneseForConditionalGeneration(config=lowercase__       )
    model.to(lowercase__       )
    model.eval()
    __lowerCamelCase	:	Optional[int]   	=			model(
        input_ids=lowercase__      ,     )
    self.parent.assertIsNotNone(lowercase__       )






@require_torch
class  lowerCamelCase_   (    _UpperCAmelCase       ,       _UpperCAmelCase       ,       unittest.TestCase					):
  """simple docstring"""




  a_							 =(GPTSanJapaneseModel,) if is_torch_available() else ()
  a_							 =(
      {
          """conversational""": GPTSanJapaneseForConditionalGeneration,
          """feature-extraction""": GPTSanJapaneseForConditionalGeneration,
          """summarization""": GPTSanJapaneseForConditionalGeneration,
          """text2text-generation""": GPTSanJapaneseForConditionalGeneration,
          """translation""": GPTSanJapaneseForConditionalGeneration,
      }
      if is_torch_available()
      else {}
  )
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
  a_							 =[0.8, 0.9]




  def 						_lowercase					(			self				:     List[Any]      ,     _a				:     List[Any]      ,     _a				:     Optional[int]      ,     _a				:     Any      ,     _a				:     Any      ,     _a				:     Dict       ) -> int:






    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 						_lowercase					(			self				:     int       ) -> List[str]:
    __lowerCamelCase	:	Dict   	=			GPTSanJapaneseTester(self       )
    __lowerCamelCase	:	Optional[Any]   	=			ConfigTester(self      ,     config_class=lowercase__      ,     d_model=37       )




  def 						_lowercase					(			self				:     str       ) -> Optional[int]:
    GPTSanJapaneseConfig()




  def 						_lowercase					(			self				:     Dict       ) -> List[str]:
    __lowerCamelCase	:	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 						_lowercase					(			self				:     List[Any]       ) -> List[Any]:
    super().test_model_parallelism()




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






@require_torch
class  lowerCamelCase_   (    _UpperCAmelCase       ,       _UpperCAmelCase       ,       unittest.TestCase					):
  """simple docstring"""




  a_							 =(GPTSanJapaneseForConditionalGeneration,) if is_torch_available() else ()
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  a_							 =False
  # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
  a_							 =[0.8, 0.9]




  def 						_lowercase					(			self				:     List[Any]       ) -> List[Any]:
    __lowerCamelCase	:	int   	=			GPTSanJapaneseTester(self       )
    __lowerCamelCase	:	List[str]   	=			ConfigTester(self      ,     config_class=lowercase__      ,     d_model=37       )




  def 						_lowercase					(			self				:     Optional[Any]       ) -> Any:
    GPTSanJapaneseConfig()




  def 						_lowercase					(			self				:     int       ) -> Optional[Any]:
    __lowerCamelCase	:	List[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 						_lowercase					(			self				:     Any       ) -> List[str]:
    super().test_model_parallelism()




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




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

    def check(_a				:     Optional[int]      ,     _a				:     Optional[int]      ,     _a				:     Optional[int]=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 						_lowercase					(			self				:     Tuple       ) -> str:
    __lowerCamelCase	:	Optional[Any]   	=			GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'       )
    __lowerCamelCase	:	Tuple   	=			GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'       )
    model.to(lowercase__       )

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

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

    __lowerCamelCase	:	List[Any]   	=			"left"
    __lowerCamelCase	:	Optional[int]   	=			tokenizer(lowercase__      ,     return_tensors='pt'      ,     padding=lowercase__       )
    __lowerCamelCase	:	Union[str, Any]   	=			inputs["input_ids"].to(lowercase__       )

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

    __lowerCamelCase	:	Any   	=			model.generate(
        input_ids=lowercase__      ,     attention_mask=inputs['attention_mask'].to(lowercase__       )      ,     max_new_tokens=3      ,     generation_config=lowercase__      ,     )

    __lowerCamelCase	:	str   	=			tokenizer(sentences[0]      ,     return_tensors='pt'       ).input_ids.to(lowercase__       )
    __lowerCamelCase	:	List[Any]   	=			model.generate(
        input_ids=lowercase__      ,     max_new_tokens=3      ,     generation_config=lowercase__       )

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

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

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

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




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




  @slow
  def 						_lowercase					(			self				:     List[Any]       ) -> str:
    __lowerCamelCase	:	Any   	=			GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'       )
    __lowerCamelCase	:	Tuple   	=			GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'       )
    model.to(lowercase__       )

    # set deterministically
    __lowerCamelCase	:	Union[str, Any]   	=			GenerationConfig.from_pretrained('Tanrei/GPTSAN-japanese'       )
    __lowerCamelCase	:	Dict   	=			1

    __lowerCamelCase	:	Tuple   	=			"武田信玄は、"
    __lowerCamelCase	:	int   	=			tokenizer(lowercase__      ,     return_tensors='pt'       ).input_ids.to(lowercase__       )
    __lowerCamelCase	:	Any   	=			tokenizer([input_text, input_text]      ,     return_tensors='pt'       ).input_ids.to(lowercase__       )

    # spout from uniform and one-hot
    __lowerCamelCase	:	Any   	=			[
        # 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
    ]

    __lowerCamelCase	:	int   	=			model.generate(
        input_ids=lowercase__      ,     spout=spouts[0]      ,     max_new_tokens=20      ,     generation_config=lowercase__      ,     )

    __lowerCamelCase	:	Optional[int]   	=			model.generate(
        input_ids=lowercase__      ,     spout=spouts[1]      ,     max_new_tokens=20      ,     generation_config=lowercase__      ,     )

    __lowerCamelCase	:	Dict   	=			model.generate(
        input_ids=lowercase__      ,     spout=lowercase__      ,     max_new_tokens=20      ,     generation_config=lowercase__      ,     )

    __lowerCamelCase	:	str   	=			tokenizer.decode(outputa[0]       )
    __lowerCamelCase	:	Any   	=			tokenizer.decode(outputa[0]       )
    __lowerCamelCase	:	Union[str, Any]   	=			tokenizer.batch_decode(lowercase__       )

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




  @slow
  def 						_lowercase					(			self				:     Tuple       ) -> Tuple:
    __lowerCamelCase	:	Optional[int]   	=			GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'       )
    __lowerCamelCase	:	List[str]   	=			GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'       )
    model.to(lowercase__       )

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

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

    __lowerCamelCase	:	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=20      ,     generation_config=lowercase__      ,     )

    __lowerCamelCase	:	Optional[int]   	=			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=20      ,     generation_config=lowercase__      ,     )

    __lowerCamelCase	:	Optional[int]   	=			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=20      ,     generation_config=lowercase__      ,     )

    __lowerCamelCase	:	List[Any]   	=			tokenizer.decode(outputa[0]       )
    __lowerCamelCase	:	str   	=			tokenizer.decode(outputa[0]       )
    __lowerCamelCase	:	Optional[int]   	=			tokenizer.batch_decode(lowercase__       )

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











'''simple docstring'''





import math
import warnings
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin






@dataclass
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete
class  lowerCamelCase_   (    SCREAMING_SNAKE_CASE_					):
  """simple docstring"""




  a_							 =42
  a_							 =None




def 	a_					(							_lowerCAmelCase	,_lowerCAmelCase=0.999	,_lowerCAmelCase="cosine"	,)				->				Union[str, Any]:
  if alpha_transform_type == "cosine":

    def alpha_bar_fn(_lowerCAmelCase	):
      return math.cos((t + 0.008) / 1.008 * math.pi / 2	) ** 2

  elif alpha_transform_type == "exp":

    def alpha_bar_fn(_lowerCAmelCase	):
      return math.exp(t * -12.0	)

  else:
    raise ValueError(F'Unsupported alpha_tranform_type: {alpha_transform_type}'	)

  __lowerCamelCase	:	Union[str, Any]   	=			[]
  for i in range(_lowerCAmelCase	):
    __lowerCamelCase	:	Dict   	=			i / num_diffusion_timesteps
    __lowerCamelCase	:	Optional[int]   	=			(i + 1) / num_diffusion_timesteps
    betas.append(min(1 - alpha_bar_fn(_lowerCAmelCase	) / alpha_bar_fn(_lowerCAmelCase	)	,_lowerCAmelCase	)	)
  return torch.tensor(_lowerCAmelCase	,dtype=torch.floataa	)






class  lowerCamelCase_   (    SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_					):
  """simple docstring"""




  a_							 =[e.name for e in KarrasDiffusionSchedulers]
  a_							 =1




  @register_to_config
  def __init__(			self				:     List[Any]      ,     _a				:     int = 1000      ,     _a				:     float = 0.0001      ,     _a				:     float = 0.02      ,     _a				:     str = "linear"      ,     _a				:     Optional[Union[np.ndarray, List[float]]] = None      ,     _a				:     Optional[bool] = False      ,     _a				:     str = "epsilon"      ,     _a				:     str = "linspace"      ,     _a				:     int = 0      ,     ) -> Tuple:
    if trained_betas is not None:
      __lowerCamelCase	:	Optional[int]   	=			torch.tensor(_a      ,     dtype=torch.floataa       )
    elif beta_schedule == "linear":
      __lowerCamelCase	:	List[str]   	=			torch.linspace(_a      ,     _a      ,     _a      ,     dtype=torch.floataa       )
    elif beta_schedule == "scaled_linear":
      # this schedule is very specific to the latent diffusion model.
      __lowerCamelCase	:	Optional[int]   	=			(
          torch.linspace(beta_start**0.5      ,     beta_end**0.5      ,     _a      ,     dtype=torch.floataa       ) ** 2
      )
    elif beta_schedule == "squaredcos_cap_v2":
      # Glide cosine schedule
      __lowerCamelCase	:	Any   	=			betas_for_alpha_bar(_a       )
    else:
      raise NotImplementedError(f'{beta_schedule} does is not implemented for {self.__class__}'       )

    __lowerCamelCase	:	Optional[Any]   	=			1.0 - self.betas
    __lowerCamelCase	:	Optional[int]   	=			torch.cumprod(self.alphas      ,     dim=0       )

    __lowerCamelCase	:	str   	=			np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5       )
    __lowerCamelCase	:	Optional[Any]   	=			np.concatenate([sigmas[::-1], [0.0]]       ).astype(np.floataa       )
    __lowerCamelCase	:	Dict   	=			torch.from_numpy(_a       )

    # setable values
    __lowerCamelCase	:	List[str]   	=			None
    __lowerCamelCase	:	int   	=			use_karras_sigmas
    self.set_timesteps(_a      ,     _a       )
    __lowerCamelCase	:	int   	=			[]
    __lowerCamelCase	:	Dict   	=			False




  @property
  def 						_lowercase					(			self				:     Optional[Any]       ) -> Any:
    # standard deviation of the initial noise distribution
    if self.config.timestep_spacing in ["linspace", "trailing"]:
      return self.sigmas.max()

    return (self.sigmas.max() ** 2 + 1) ** 0.5




  def 						_lowercase					(			self				:     List[Any]      ,     _a				:     torch.FloatTensor      ,     _a				:     Union[float, torch.FloatTensor]       ) -> torch.FloatTensor:
    if isinstance(_a      ,     torch.Tensor       ):
      __lowerCamelCase	:	List[Any]   	=			timestep.to(self.timesteps.device       )
    __lowerCamelCase	:	int   	=			(self.timesteps == timestep).nonzero().item()
    __lowerCamelCase	:	Any   	=			self.sigmas[step_index]
    __lowerCamelCase	:	List[str]   	=			sample / ((sigma**2 + 1) ** 0.5)
    __lowerCamelCase	:	Any   	=			True
    return sample




  def 						_lowercase					(			self				:     Optional[int]      ,     _a				:     str      ,     _a				:     List[str]      ,     _a				:     Union[str, Any]       ) -> Optional[Any]:

    def lms_derivative(_a				:     List[str]       ):
      __lowerCamelCase	:	Any   	=			1.0
      for k in range(_a       ):
        if current_order == k:
          continue
        prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k])
      return prod

    __lowerCamelCase	:	Any   	=			integrate.quad(_a      ,     self.sigmas[t]      ,     self.sigmas[t + 1]      ,     epsrel=1e-4       )[0]

    return integrated_coeff




  def 						_lowercase					(			self				:     int      ,     _a				:     int      ,     _a				:     Union[str, torch.device] = None       ) -> str:
    __lowerCamelCase	:	Optional[Any]   	=			num_inference_steps

    # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
    if self.config.timestep_spacing == "linspace":
      __lowerCamelCase	:	Any   	=			np.linspace(0      ,     self.config.num_train_timesteps - 1      ,     _a      ,     dtype=_a       )[
          ::-1
      ].copy()
    elif self.config.timestep_spacing == "leading":
      __lowerCamelCase	:	Union[str, Any]   	=			self.config.num_train_timesteps // self.num_inference_steps
      # creates integer timesteps by multiplying by ratio
      # casting to int to avoid issues when num_inference_step is power of 3
      __lowerCamelCase	:	Union[str, Any]   	=			(np.arange(0      ,     _a       ) * step_ratio).round()[::-1].copy().astype(_a       )
      timesteps += self.config.steps_offset
    elif self.config.timestep_spacing == "trailing":
      __lowerCamelCase	:	int   	=			self.config.num_train_timesteps / self.num_inference_steps
      # creates integer timesteps by multiplying by ratio
      # casting to int to avoid issues when num_inference_step is power of 3
      __lowerCamelCase	:	str   	=			(np.arange(self.config.num_train_timesteps      ,     0      ,     -step_ratio       )).round().copy().astype(_a       )
      timesteps -= 1
    else:
      raise ValueError(
          f'{self.config.timestep_spacing} is not supported. Please make sure to choose one of \'linspace\', \'leading\' or \'trailing\'.'       )

    __lowerCamelCase	:	Optional[Any]   	=			np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5       )
    __lowerCamelCase	:	List[Any]   	=			np.log(_a       )
    __lowerCamelCase	:	Optional[int]   	=			np.interp(_a      ,     np.arange(0      ,     len(_a       )       )      ,     _a       )

    if self.use_karras_sigmas:
      __lowerCamelCase	:	int   	=			self._convert_to_karras(in_sigmas=_a       )
      __lowerCamelCase	:	str   	=			np.array([self._sigma_to_t(_a      ,     _a       ) for sigma in sigmas]       )

    __lowerCamelCase	:	Optional[Any]   	=			np.concatenate([sigmas, [0.0]]       ).astype(np.floataa       )

    __lowerCamelCase	:	Dict   	=			torch.from_numpy(_a       ).to(device=_a       )
    if str(_a       ).startswith('mps'       ):
      # mps does not support float64
      __lowerCamelCase	:	List[str]   	=			torch.from_numpy(_a       ).to(_a      ,     dtype=torch.floataa       )
    else:
      __lowerCamelCase	:	str   	=			torch.from_numpy(_a       ).to(device=_a       )

    __lowerCamelCase	:	List[Any]   	=			[]




  def 						_lowercase					(			self				:     List[Any]      ,     _a				:     List[Any]      ,     _a				:     int       ) -> Tuple:
    # get log sigma
    __lowerCamelCase	:	Optional[int]   	=			np.log(_a       )

    # get distribution
    __lowerCamelCase	:	int   	=			log_sigma - log_sigmas[:, np.newaxis]

    # get sigmas range
    __lowerCamelCase	:	Dict   	=			np.cumsum((dists >= 0)      ,     axis=0       ).argmax(axis=0       ).clip(max=log_sigmas.shape[0] - 2       )
    __lowerCamelCase	:	List[Any]   	=			low_idx + 1

    __lowerCamelCase	:	Optional[Any]   	=			log_sigmas[low_idx]
    __lowerCamelCase	:	str   	=			log_sigmas[high_idx]

    # interpolate sigmas
    __lowerCamelCase	:	List[str]   	=			(low - log_sigma) / (low - high)
    __lowerCamelCase	:	List[str]   	=			np.clip(_a      ,     0      ,     1       )

    # transform interpolation to time range
    __lowerCamelCase	:	Dict   	=			(1 - w) * low_idx + w * high_idx
    __lowerCamelCase	:	List[str]   	=			t.reshape(sigma.shape       )
    return t




  def 						_lowercase					(			self				:     int      ,     _a				:     torch.FloatTensor       ) -> torch.FloatTensor:
    __lowerCamelCase	:	float   	=			in_sigmas[-1].item()
    __lowerCamelCase	:	float   	=			in_sigmas[0].item()

    __lowerCamelCase	:	Optional[Any]   	=			7.0  # 7.0 is the value used in the paper
    __lowerCamelCase	:	Any   	=			np.linspace(0      ,     1      ,     self.num_inference_steps       )
    __lowerCamelCase	:	Optional[Any]   	=			sigma_min ** (1 / rho)
    __lowerCamelCase	:	Tuple   	=			sigma_max ** (1 / rho)
    __lowerCamelCase	:	str   	=			(max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
    return sigmas




  def 						_lowercase					(			self				:     Dict      ,     _a				:     torch.FloatTensor      ,     _a				:     Union[float, torch.FloatTensor]      ,     _a				:     torch.FloatTensor      ,     _a				:     int = 4      ,     _a				:     bool = True      ,     ) -> Union[LMSDiscreteSchedulerOutput, Tuple]:
    if not self.is_scale_input_called:
      warnings.warn(
          'The `scale_model_input` function should be called before `step` to ensure correct denoising. '
          'See `StableDiffusionPipeline` for a usage example.'       )

    if isinstance(_a      ,     torch.Tensor       ):
      __lowerCamelCase	:	Optional[int]   	=			timestep.to(self.timesteps.device       )
    __lowerCamelCase	:	List[str]   	=			(self.timesteps == timestep).nonzero().item()
    __lowerCamelCase	:	Optional[int]   	=			self.sigmas[step_index]

    # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
    if self.config.prediction_type == "epsilon":
      __lowerCamelCase	:	Optional[Any]   	=			sample - sigma * model_output
    elif self.config.prediction_type == "v_prediction":
      # * c_out + input * c_skip
      __lowerCamelCase	:	List[str]   	=			model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
    elif self.config.prediction_type == "sample":
      __lowerCamelCase	:	Union[str, Any]   	=			model_output
    else:
      raise ValueError(
          f'prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`'       )

    # 2. Convert to an ODE derivative
    __lowerCamelCase	:	str   	=			(sample - pred_original_sample) / sigma
    self.derivatives.append(_a       )
    if len(self.derivatives       ) > order:
      self.derivatives.pop(0       )

    # 3. Compute linear multistep coefficients
    __lowerCamelCase	:	List[str]   	=			min(step_index + 1      ,     _a       )
    __lowerCamelCase	:	int   	=			[self.get_lms_coefficient(_a      ,     _a      ,     _a       ) for curr_order in range(_a       )]

    # 4. Compute previous sample based on the derivatives path
    __lowerCamelCase	:	Optional[Any]   	=			sample + sum(
        coeff * derivative for coeff, derivative in zip(_a      ,     reversed(self.derivatives       )       )       )

    if not return_dict:
      return (prev_sample,)

    return LMSDiscreteSchedulerOutput(prev_sample=_a      ,     pred_original_sample=_a       )




  def 						_lowercase					(			self				:     List[str]      ,     _a				:     torch.FloatTensor      ,     _a				:     torch.FloatTensor      ,     _a				:     torch.FloatTensor      ,     ) -> torch.FloatTensor:
    # Make sure sigmas and timesteps have the same device and dtype as original_samples
    __lowerCamelCase	:	Any   	=			self.sigmas.to(device=original_samples.device      ,     dtype=original_samples.dtype       )
    if original_samples.device.type == "mps" and torch.is_floating_point(_a       ):
      # mps does not support float64
      __lowerCamelCase	:	Optional[int]   	=			self.timesteps.to(original_samples.device      ,     dtype=torch.floataa       )
      __lowerCamelCase	:	Optional[int]   	=			timesteps.to(original_samples.device      ,     dtype=torch.floataa       )
    else:
      __lowerCamelCase	:	int   	=			self.timesteps.to(original_samples.device       )
      __lowerCamelCase	:	Optional[Any]   	=			timesteps.to(original_samples.device       )

    __lowerCamelCase	:	List[str]   	=			[(schedule_timesteps == t).nonzero().item() for t in timesteps]

    __lowerCamelCase	:	Union[str, Any]   	=			sigmas[step_indices].flatten()
    while len(sigma.shape       ) < len(original_samples.shape       ):
      __lowerCamelCase	:	List[Any]   	=			sigma.unsqueeze(-1       )

    __lowerCamelCase	:	Optional[Any]   	=			original_samples + noise * sigma
    return noisy_samples




  def __len__(			self				:     Union[str, Any]       ) -> int:
    return self.config.num_train_timesteps






