# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()







class __snake_case				(						lowerCamelCase_   ,		lowerCamelCase_   ,		lowerCamelCase_   ,		unittest.TestCase		):
   lowerCAmelCase_														=							StableDiffusionControlNetImgaImgPipeline
   lowerCAmelCase_														=							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
   lowerCAmelCase_														=							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   lowerCAmelCase_														=							IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}		)
   lowerCAmelCase_														=							IMAGE_TO_IMAGE_IMAGE_PARAMS




   def 	__a						(    self			:				List[str]      ):



     """simple docstring"""






     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      UNetaDConditionModel(
         block_out_channels=(32, 64)       ,	layers_per_block=2       ,	sample_size=32       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,	up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")       ,	cross_attention_dim=32       ,	)
     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      ControlNetModel(
         block_out_channels=(32, 64)       ,	layers_per_block=2       ,	in_channels=4       ,	down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,	cross_attention_dim=32       ,	conditioning_embedding_out_channels=(16, 32)       ,	)
     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      DDIMScheduler(
         beta_start=0.0_00_85       ,	beta_end=0.0_12       ,	beta_schedule="""scaled_linear"""       ,	clip_sample=_lowercase       ,	set_alpha_to_one=_lowercase       ,	)
     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      AutoencoderKL(
         block_out_channels=[32, 64]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]       ,	up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]       ,	latent_channels=4       ,	)
     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      CLIPTextConfig(
         bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=32       ,	intermediate_size=37       ,	layer_norm_eps=1E-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=10_00       ,	)
     SCREAMING_SNAKE_CASE__										=      CLIPTextModel(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

     SCREAMING_SNAKE_CASE__										=      {
         """unet""": unet,
         """controlnet""": controlnet,
         """scheduler""": scheduler,
         """vae""": vae,
         """text_encoder""": text_encoder,
         """tokenizer""": tokenizer,
         """safety_checker""": None,
         """feature_extractor""": None,
     }
     return components




   def 	__a						(    self			:				str       ,	_lowercase			:				Tuple       ,	_lowercase			:				Any=0      ):



     """simple docstring"""






     if str(_lowercase      ).startswith("""mps"""      ):
       SCREAMING_SNAKE_CASE__										=      torch.manual_seed(_lowercase      )
     else:
       SCREAMING_SNAKE_CASE__										=      torch.Generator(device=_lowercase      ).manual_seed(_lowercase      )

     SCREAMING_SNAKE_CASE__										=      2
     SCREAMING_SNAKE_CASE__										=      randn_tensor(
         (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,	generator=_lowercase       ,	device=torch.device(_lowercase      )       ,	)
     SCREAMING_SNAKE_CASE__										=      floats_tensor(control_image.shape       ,	rng=random.Random(_lowercase      )      ).to(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      image.cpu().permute(0       ,	2       ,	3       ,	1      )[0]
     SCREAMING_SNAKE_CASE__										=      Image.fromarray(np.uinta(_lowercase      )      ).convert("""RGB"""      ).resize((64, 64)      )
     SCREAMING_SNAKE_CASE__										=      {
         """prompt""": """A painting of a squirrel eating a burger""",
         """generator""": generator,
         """num_inference_steps""": 2,
         """guidance_scale""": 6.0,
         """output_type""": """numpy""",
         """image""": image,
         """control_image""": control_image,
     }

     return inputs




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



     """simple docstring"""






     return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3      )




   @unittest.skipIf(
       torch_device != """cuda""" or not is_xformers_available()       ,	reason="""XFormers attention is only available with CUDA and `xformers` installed"""       ,	)
   def 	__a						(    self			:				Tuple      ):



     """simple docstring"""






     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3      )




   def 	__a						(    self			:				List[str]      ):



     """simple docstring"""






     self._test_inference_batch_single_identical(expected_max_diff=2E-3      )







class __snake_case				(						lowerCamelCase_   ,		lowerCamelCase_   ,		unittest.TestCase		):
   lowerCAmelCase_														=							StableDiffusionControlNetImgaImgPipeline
   lowerCAmelCase_														=							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
   lowerCAmelCase_														=							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   lowerCAmelCase_														=							frozenset([]		)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess




   def 	__a						(    self			:				List[str]      ):



     """simple docstring"""






     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      UNetaDConditionModel(
         block_out_channels=(32, 64)       ,	layers_per_block=2       ,	sample_size=32       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,	up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")       ,	cross_attention_dim=32       ,	)
     torch.manual_seed(0      )

     def init_weights(_lowercase			:				List[Any]      ):
       if isinstance(_lowercase       ,	torch.nn.Convad      ):
         torch.nn.init.normal(m.weight      )
         m.bias.data.fill_(1.0      )

     SCREAMING_SNAKE_CASE__										=      ControlNetModel(
         block_out_channels=(32, 64)       ,	layers_per_block=2       ,	in_channels=4       ,	down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,	cross_attention_dim=32       ,	conditioning_embedding_out_channels=(16, 32)       ,	)
     controlneta.controlnet_down_blocks.apply(_lowercase      )

     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      ControlNetModel(
         block_out_channels=(32, 64)       ,	layers_per_block=2       ,	in_channels=4       ,	down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,	cross_attention_dim=32       ,	conditioning_embedding_out_channels=(16, 32)       ,	)
     controlneta.controlnet_down_blocks.apply(_lowercase      )

     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      DDIMScheduler(
         beta_start=0.0_00_85       ,	beta_end=0.0_12       ,	beta_schedule="""scaled_linear"""       ,	clip_sample=_lowercase       ,	set_alpha_to_one=_lowercase       ,	)
     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      AutoencoderKL(
         block_out_channels=[32, 64]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]       ,	up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]       ,	latent_channels=4       ,	)
     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      CLIPTextConfig(
         bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=32       ,	intermediate_size=37       ,	layer_norm_eps=1E-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=10_00       ,	)
     SCREAMING_SNAKE_CASE__										=      CLIPTextModel(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

     SCREAMING_SNAKE_CASE__										=      MultiControlNetModel([controlneta, controlneta]      )

     SCREAMING_SNAKE_CASE__										=      {
         """unet""": unet,
         """controlnet""": controlnet,
         """scheduler""": scheduler,
         """vae""": vae,
         """text_encoder""": text_encoder,
         """tokenizer""": tokenizer,
         """safety_checker""": None,
         """feature_extractor""": None,
     }
     return components




   def 	__a						(    self			:				Any       ,	_lowercase			:				str       ,	_lowercase			:				str=0      ):



     """simple docstring"""






     if str(_lowercase      ).startswith("""mps"""      ):
       SCREAMING_SNAKE_CASE__										=      torch.manual_seed(_lowercase      )
     else:
       SCREAMING_SNAKE_CASE__										=      torch.Generator(device=_lowercase      ).manual_seed(_lowercase      )

     SCREAMING_SNAKE_CASE__										=      2

     SCREAMING_SNAKE_CASE__										=      [
         randn_tensor(
             (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,	generator=_lowercase       ,	device=torch.device(_lowercase      )       ,	),
         randn_tensor(
             (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,	generator=_lowercase       ,	device=torch.device(_lowercase      )       ,	),
     ]

     SCREAMING_SNAKE_CASE__										=      floats_tensor(control_image[0].shape       ,	rng=random.Random(_lowercase      )      ).to(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      image.cpu().permute(0       ,	2       ,	3       ,	1      )[0]
     SCREAMING_SNAKE_CASE__										=      Image.fromarray(np.uinta(_lowercase      )      ).convert("""RGB"""      ).resize((64, 64)      )
     SCREAMING_SNAKE_CASE__										=      {
         """prompt""": """A painting of a squirrel eating a burger""",
         """generator""": generator,
         """num_inference_steps""": 2,
         """guidance_scale""": 6.0,
         """output_type""": """numpy""",
         """image""": image,
         """control_image""": control_image,
     }

     return inputs




   def 	__a						(    self			:				Optional[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.get_dummy_components()
     SCREAMING_SNAKE_CASE__										=      self.pipeline_class(**_lowercase      )
     pipe.to(_lowercase      )

     SCREAMING_SNAKE_CASE__										=      10.0
     SCREAMING_SNAKE_CASE__										=      4

     SCREAMING_SNAKE_CASE__										=      self.get_dummy_inputs(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      steps
     SCREAMING_SNAKE_CASE__										=      scale
     SCREAMING_SNAKE_CASE__										=      pipe(**_lowercase      )[0]

     SCREAMING_SNAKE_CASE__										=      self.get_dummy_inputs(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      steps
     SCREAMING_SNAKE_CASE__										=      scale
     SCREAMING_SNAKE_CASE__										=      pipe(**_lowercase       ,	control_guidance_start=0.1       ,	control_guidance_end=0.2      )[0]

     SCREAMING_SNAKE_CASE__										=      self.get_dummy_inputs(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      steps
     SCREAMING_SNAKE_CASE__										=      scale
     SCREAMING_SNAKE_CASE__										=      pipe(**_lowercase       ,	control_guidance_start=[0.1, 0.3]       ,	control_guidance_end=[0.2, 0.7]      )[0]

     SCREAMING_SNAKE_CASE__										=      self.get_dummy_inputs(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      steps
     SCREAMING_SNAKE_CASE__										=      scale
     SCREAMING_SNAKE_CASE__										=      pipe(**_lowercase       ,	control_guidance_start=0.4       ,	control_guidance_end=[0.5, 0.8]      )[0]

     # make sure that all outputs are different
     assert np.sum(np.abs(output_a - output_a      )      ) > 1E-3
     assert np.sum(np.abs(output_a - output_a      )      ) > 1E-3
     assert np.sum(np.abs(output_a - output_a      )      ) > 1E-3




   def 	__a						(    self			:				List[Any]      ):



     """simple docstring"""






     return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3      )




   @unittest.skipIf(
       torch_device != """cuda""" or not is_xformers_available()       ,	reason="""XFormers attention is only available with CUDA and `xformers` installed"""       ,	)
   def 	__a						(    self			:				Optional[Any]      ):



     """simple docstring"""






     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3      )




   def 	__a						(    self			:				Dict      ):



     """simple docstring"""






     self._test_inference_batch_single_identical(expected_max_diff=2E-3      )




   def 	__a						(    self			:				int      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.get_dummy_components()
     SCREAMING_SNAKE_CASE__										=      self.pipeline_class(**_lowercase      )
     pipe.to(_lowercase      )
     pipe.set_progress_bar_config(disable=_lowercase      )
     with tempfile.TemporaryDirectory() as tmpdir:
       try:
         # save_pretrained is not implemented for Multi-ControlNet
         pipe.save_pretrained(_lowercase      )
       except NotImplementedError:
         pass







@slow
@require_torch_gpu
class __snake_case				(						unittest.TestCase		):




   def 	__a						(    self			:				str      ):



     """simple docstring"""






     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




   def 	__a						(    self			:				Optional[int]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""      )

     SCREAMING_SNAKE_CASE__										=      StableDiffusionControlNetImgaImgPipeline.from_pretrained(
         """runwayml/stable-diffusion-v1-5"""       ,	safety_checker=_lowercase       ,	controlnet=_lowercase      )
     pipe.enable_model_cpu_offload()
     pipe.set_progress_bar_config(disable=_lowercase      )

     SCREAMING_SNAKE_CASE__										=      torch.Generator(device="""cpu"""      ).manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      """evil space-punk bird"""
     SCREAMING_SNAKE_CASE__										=      load_image(
         """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"""      ).resize((5_12, 5_12)      )
     SCREAMING_SNAKE_CASE__										=      load_image(
         """https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"""      ).resize((5_12, 5_12)      )

     SCREAMING_SNAKE_CASE__										=      pipe(
         _lowercase       ,	_lowercase       ,	control_image=_lowercase       ,	generator=_lowercase       ,	output_type="""np"""       ,	num_inference_steps=50       ,	strength=0.6       ,	)

     SCREAMING_SNAKE_CASE__										=      output.images[0]

     assert image.shape == (5_12, 5_12, 3)

     SCREAMING_SNAKE_CASE__										=      load_numpy(
         """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy"""      )

     assert np.abs(expected_image - image      ).max() < 9E-2

import unittest

from transformers import GPTNeoXJapaneseConfig, is_torch_available
from transformers.models.gpt_neox_japanese.tokenization_gpt_neox_japanese import GPTNeoXJapaneseTokenizer
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import GPTNeoXJapaneseForCausalLM, GPTNeoXJapaneseModel







class __snake_case				:




   def __init__(    self			:				Optional[int]       ,	_lowercase			:				Union[str, Any]       ,	_lowercase			:				List[Any]=13       ,	_lowercase			:				List[Any]=7       ,	_lowercase			:				Optional[int]=True       ,	_lowercase			:				str=True       ,	_lowercase			:				str=True       ,	_lowercase			:				List[str]=True       ,	_lowercase			:				List[str]=99       ,	_lowercase			:				List[str]=32       ,	_lowercase			:				str=5       ,	_lowercase			:				str=4       ,	_lowercase			:				str=4       ,	_lowercase			:				Union[str, Any]="gelu"       ,	_lowercase			:				str=0.0       ,	_lowercase			:				Union[str, Any]=0.1       ,	_lowercase			:				List[str]=True       ,	_lowercase			:				Union[str, Any]=5_12       ,	_lowercase			:				List[str]=16       ,	_lowercase			:				Dict=2       ,	_lowercase			:				int=0.02       ,	_lowercase			:				Any=3       ,	_lowercase			:				int=4       ,	_lowercase			:				List[str]=None       ,	):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      parent
     SCREAMING_SNAKE_CASE__										=      batch_size
     SCREAMING_SNAKE_CASE__										=      seq_length
     SCREAMING_SNAKE_CASE__										=      is_training
     SCREAMING_SNAKE_CASE__										=      use_input_mask
     SCREAMING_SNAKE_CASE__										=      use_token_type_ids
     SCREAMING_SNAKE_CASE__										=      use_labels
     SCREAMING_SNAKE_CASE__										=      vocab_size
     SCREAMING_SNAKE_CASE__										=      hidden_size
     SCREAMING_SNAKE_CASE__										=      num_hidden_layers
     SCREAMING_SNAKE_CASE__										=      num_attention_heads
     SCREAMING_SNAKE_CASE__										=      intermediate_multiple_size
     SCREAMING_SNAKE_CASE__										=      hidden_act
     SCREAMING_SNAKE_CASE__										=      hidden_dropout
     SCREAMING_SNAKE_CASE__										=      attention_dropout
     SCREAMING_SNAKE_CASE__										=      weight_tying
     SCREAMING_SNAKE_CASE__										=      max_position_embeddings
     SCREAMING_SNAKE_CASE__										=      type_vocab_size
     SCREAMING_SNAKE_CASE__										=      type_sequence_label_size
     SCREAMING_SNAKE_CASE__										=      initializer_range
     SCREAMING_SNAKE_CASE__										=      num_labels
     SCREAMING_SNAKE_CASE__										=      num_choices
     SCREAMING_SNAKE_CASE__										=      scope




   def 	__a						(    self			:				List[Any]      ):



     """simple docstring"""






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

     SCREAMING_SNAKE_CASE__										=      None
     if self.use_input_mask:
       SCREAMING_SNAKE_CASE__										=      random_attention_mask([self.batch_size, self.seq_length]      )

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

     SCREAMING_SNAKE_CASE__										=      self.get_config()

     return config, input_ids, input_mask, token_labels




   def 	__a						(    self			:				Optional[int]      ):



     """simple docstring"""






     return GPTNeoXJapaneseConfig(
         vocab_size=self.vocab_size       ,	hidden_size=self.hidden_size       ,	num_hidden_layers=self.num_hidden_layers       ,	num_attention_heads=self.num_attention_heads       ,	intermediate_multiple_size=self.intermediate_multiple_size       ,	hidden_act=self.hidden_act       ,	hidden_dropout=self.hidden_dropout       ,	attention_dropout=self.attention_dropout       ,	weight_tying=self.weight_tying       ,	max_position_embeddings=self.max_position_embeddings       ,	type_vocab_size=self.type_vocab_size       ,	is_decoder=_lowercase       ,	initializer_range=self.initializer_range       ,	)




   def 	__a						(    self			:				Optional[int]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.prepare_config_and_inputs()

     SCREAMING_SNAKE_CASE__										=      True

     return config, input_ids, input_mask, token_labels




   def 	__a						(    self			:				List[Any]       ,	_lowercase			:				Optional[Any]       ,	_lowercase			:				int       ,	_lowercase			:				List[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      GPTNeoXJapaneseModel(config=_lowercase      )
     model.to(_lowercase      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	attention_mask=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      model(_lowercase      )
     self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)      )




   def 	__a						(    self			:				Dict       ,	_lowercase			:				int       ,	_lowercase			:				str       ,	_lowercase			:				Any      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      True
     SCREAMING_SNAKE_CASE__										=      GPTNeoXJapaneseModel(_lowercase      )
     model.to(_lowercase      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	attention_mask=_lowercase      )
     self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)      )




   def 	__a						(    self			:				Any       ,	_lowercase			:				Dict       ,	_lowercase			:				List[Any]       ,	_lowercase			:				Optional[Any]       ,	_lowercase			:				Dict      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      GPTNeoXJapaneseForCausalLM(config=_lowercase      )
     model.to(_lowercase      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	attention_mask=_lowercase       ,	labels=_lowercase      )
     self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)      )




   def 	__a						(    self			:				int       ,	_lowercase			:				Dict       ,	_lowercase			:				Optional[Any]       ,	_lowercase			:				int      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      True
     SCREAMING_SNAKE_CASE__										=      GPTNeoXJapaneseForCausalLM(config=_lowercase      )
     model.to(_lowercase      )
     model.eval()

     # first forward pass
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	attention_mask=_lowercase       ,	use_cache=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     SCREAMING_SNAKE_CASE__										=      ids_tensor((self.batch_size, 3)       ,	config.vocab_size      )
     SCREAMING_SNAKE_CASE__										=      ids_tensor((self.batch_size, 3)       ,	vocab_size=2      )

     # append to next input_ids and
     SCREAMING_SNAKE_CASE__										=      torch.cat([input_ids, next_tokens]       ,	dim=-1      )
     SCREAMING_SNAKE_CASE__										=      torch.cat([input_mask, next_mask]       ,	dim=-1      )

     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	attention_mask=_lowercase       ,	output_hidden_states=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      output_from_no_past["""hidden_states"""][0]
     SCREAMING_SNAKE_CASE__										=      model(
         _lowercase       ,	attention_mask=_lowercase       ,	past_key_values=_lowercase       ,	output_hidden_states=_lowercase       ,	)["""hidden_states"""][0]

     # select random slice
     SCREAMING_SNAKE_CASE__										=      ids_tensor((1,)       ,	output_from_past.shape[-1]      ).item()
     SCREAMING_SNAKE_CASE__										=      output_from_no_past[:, -3:, random_slice_idx].detach()
     SCREAMING_SNAKE_CASE__										=      output_from_past[:, :, random_slice_idx].detach()

     self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]      )

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(_lowercase       ,	_lowercase       ,	atol=1E-3      )      )




   def 	__a						(    self			:				str      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.prepare_config_and_inputs()
     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      config_and_inputs
     SCREAMING_SNAKE_CASE__										=      {"""input_ids""": input_ids, """attention_mask""": input_mask}
     return config, inputs_dict







@require_torch
class __snake_case				(						lowerCamelCase_   ,		lowerCamelCase_   ,		unittest.TestCase		):
   lowerCAmelCase_														=							(GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
   lowerCAmelCase_														=							(GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
   lowerCAmelCase_														=							(
       {"feature-extraction": GPTNeoXJapaneseModel, "text-generation": GPTNeoXJapaneseForCausalLM}
       if is_torch_available()
       else {}
   )
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False




   def 	__a						(    self			:				List[str]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      GPTNeoXJapaneseModelTester(self      )
     SCREAMING_SNAKE_CASE__										=      ConfigTester(self       ,	config_class=_lowercase       ,	hidden_size=37      )




   def 	__a						(    self			:				Tuple      ):



     """simple docstring"""






     self.config_tester.run_common_tests()




   def 	__a						(    self			:				Optional[int]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(_lowercase       ,	_lowercase       ,	_lowercase      )




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



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(_lowercase       ,	_lowercase       ,	_lowercase      )




   def 	__a						(    self			:				str      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs_for_decoder()

     SCREAMING_SNAKE_CASE__										=      None

     self.model_tester.create_and_check_model_as_decoder(_lowercase       ,	_lowercase       ,	_lowercase      )




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



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_decoder_model_past_large_inputs(_lowercase       ,	_lowercase       ,	_lowercase      )




   def 	__a						(    self			:				List[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_causal_lm(*_lowercase      )




   @slow
   def 	__a						(    self			:				Any      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      """abeja/gpt-neox-japanese-2.7b"""

     SCREAMING_SNAKE_CASE__										=      ["""データサイエンティストとは、""", """100年後に必要とされる会社は、""", """フルリモートの環境で働くために必要なことは、""", """国境の長いトンネルを抜けると""", """美味しい日本食といえば、"""]

     SCREAMING_SNAKE_CASE__										=      [
         """データサイエンティストとは、データを分析し、ビジネスに役立つ知見を導き出す専門家のことです。""",
         """100年後に必要とされる会社は、「人」が中心の会社です。""",
         """フルリモートの環境で働くために必要なことは、「自分の時間をコントロールする」ことです。""",
         """国境の長いトンネルを抜けると、そこは雪国だった。""",
         """美味しい日本食といえば、やっぱりお寿司ですよね。""",
     ]

     SCREAMING_SNAKE_CASE__										=      GPTNeoXJapaneseTokenizer.from_pretrained(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      GPTNeoXJapaneseForCausalLM.from_pretrained(_lowercase      )

     SCREAMING_SNAKE_CASE__										=      []
     for prompt in prompts:
       SCREAMING_SNAKE_CASE__										=      tokenizer(_lowercase       ,	return_tensors="""pt"""      ).input_ids
       SCREAMING_SNAKE_CASE__										=      model.generate(_lowercase       ,	max_length=50      )
       SCREAMING_SNAKE_CASE__										=      tokenizer.batch_decode(_lowercase       ,	skip_special_tokens=_lowercase      )
       predicted_outputs += generated_string
     self.assertListEqual(_lowercase       ,	_lowercase      )
