

'''simple docstring'''




# This model implementation is heavily based on:

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,
    StableDiffusionControlNetInpaintPipeline,
    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 (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
    TEXT_TO_IMAGE_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()



class        lowercase_	(		a ,       a ,       a ,       unittest.TestCase ):

  '''simple docstring'''





  __lowerCAmelCase					: List[str]      =   StableDiffusionControlNetInpaintPipeline
  __lowerCAmelCase					: Any      =   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
  __lowerCAmelCase					: List[Any]      =   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
  __lowerCAmelCase					: int      =   frozenset({"control_image"} )  # skip `image` and `mask` for now, only test for control_image
  __lowerCAmelCase					: List[Any]      =   TEXT_TO_IMAGE_IMAGE_PARAMS
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    torch.manual_seed(0     )
    UpperCAmelCase													=	UNetaDConditionModel(
        block_out_channels=(3_2, 6_4)   , layers_per_block=2   , sample_size=3_2   , in_channels=9   , out_channels=4   , down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   , up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   , cross_attention_dim=3_2   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	ControlNetModel(
        block_out_channels=(3_2, 6_4)   , layers_per_block=2   , in_channels=4   , down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   , cross_attention_dim=3_2   , conditioning_embedding_out_channels=(1_6, 3_2)   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	DDIMScheduler(
        beta_start=0.0_0085   , beta_end=0.012   , beta_schedule='scaled_linear'   , clip_sample=a_   , set_alpha_to_one=a_   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	AutoencoderKL(
        block_out_channels=[3_2, 6_4]   , in_channels=3   , out_channels=3   , down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   , up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   , latent_channels=4   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	CLIPTextConfig(
        bos_token_id=0   , eos_token_id=2   , hidden_size=3_2   , intermediate_size=3_7   , layer_norm_eps=1E-05   , num_attention_heads=4   , num_hidden_layers=5   , pad_token_id=1   , vocab_size=1_0_0_0   , )
    UpperCAmelCase													=	CLIPTextModel(a_     )
    UpperCAmelCase													=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

    UpperCAmelCase													=	{
        'unet': unet,
        'controlnet': controlnet,
        'scheduler': scheduler,
        'vae': vae,
        'text_encoder': text_encoder,
        'tokenizer': tokenizer,
        'safety_checker': None,
        'feature_extractor': None,
    }
    return components
  def 						snake_case_			(     self   , a_   , a_=0     )							->				Tuple:


    """simple docstring"""
    if str(a_     ).startswith('mps'     ):
      UpperCAmelCase													=	torch.manual_seed(a_     )
    else:
      UpperCAmelCase													=	torch.Generator(device=a_     ).manual_seed(a_     )

    UpperCAmelCase													=	2
    UpperCAmelCase													=	randn_tensor(
        (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)   , generator=a_   , device=torch.device(a_     )   , )
    UpperCAmelCase													=	floats_tensor((1, 3, 3_2, 3_2)   , rng=random.Random(a_     )     ).to(a_     )
    UpperCAmelCase													=	init_image.cpu().permute(0   , 2   , 3   , 1     )[0]

    UpperCAmelCase													=	Image.fromarray(np.uinta(a_     )     ).convert('RGB'     ).resize((6_4, 6_4)     )
    UpperCAmelCase													=	Image.fromarray(np.uinta(init_image + 4     )     ).convert('RGB'     ).resize((6_4, 6_4)     )

    UpperCAmelCase													=	{
        'prompt': 'A painting of a squirrel eating a burger',
        'generator': generator,
        'num_inference_steps': 2,
        'guidance_scale': 6.0,
        'output_type': 'numpy',
        'image': image,
        'mask_image': mask_image,
        'control_image': control_image,
    }

    return inputs
  def 						snake_case_			(     self     )							->				Tuple:


    """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 						snake_case_			(     self     )							->				Union[str, Any]:


    """simple docstring"""
    self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3     )





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


    """simple docstring"""
    self._test_inference_batch_single_identical(expected_max_diff=2E-3     )




class        lowercase_	(		a ):

  '''simple docstring'''





  __lowerCAmelCase					: int      =   StableDiffusionControlNetInpaintPipeline
  __lowerCAmelCase					: Optional[int]      =   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
  __lowerCAmelCase					: List[Any]      =   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
  __lowerCAmelCase					: Any      =   frozenset([] )
  def 						snake_case_			(     self     )							->				Union[str, Any]:


    """simple docstring"""
    torch.manual_seed(0     )
    UpperCAmelCase													=	UNetaDConditionModel(
        block_out_channels=(3_2, 6_4)   , layers_per_block=2   , sample_size=3_2   , in_channels=4   , out_channels=4   , down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   , up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   , cross_attention_dim=3_2   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	ControlNetModel(
        block_out_channels=(3_2, 6_4)   , layers_per_block=2   , in_channels=4   , down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   , cross_attention_dim=3_2   , conditioning_embedding_out_channels=(1_6, 3_2)   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	DDIMScheduler(
        beta_start=0.0_0085   , beta_end=0.012   , beta_schedule='scaled_linear'   , clip_sample=a_   , set_alpha_to_one=a_   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	AutoencoderKL(
        block_out_channels=[3_2, 6_4]   , in_channels=3   , out_channels=3   , down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   , up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   , latent_channels=4   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	CLIPTextConfig(
        bos_token_id=0   , eos_token_id=2   , hidden_size=3_2   , intermediate_size=3_7   , layer_norm_eps=1E-05   , num_attention_heads=4   , num_hidden_layers=5   , pad_token_id=1   , vocab_size=1_0_0_0   , )
    UpperCAmelCase													=	CLIPTextModel(a_     )
    UpperCAmelCase													=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

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




class        lowercase_	(		a ,       a ,       unittest.TestCase ):

  '''simple docstring'''





  __lowerCAmelCase					: str      =   StableDiffusionControlNetInpaintPipeline
  __lowerCAmelCase					: List[str]      =   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
  __lowerCAmelCase					: List[Any]      =   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    torch.manual_seed(0     )
    UpperCAmelCase													=	UNetaDConditionModel(
        block_out_channels=(3_2, 6_4)   , layers_per_block=2   , sample_size=3_2   , in_channels=9   , out_channels=4   , down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   , up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   , cross_attention_dim=3_2   , )
    torch.manual_seed(0     )

    def init_weights(a_     ):
      if isinstance(a_   , torch.nn.Convad     ):
        torch.nn.init.normal(m.weight     )
        m.bias.data.fill_(1.0     )

    UpperCAmelCase													=	ControlNetModel(
        block_out_channels=(3_2, 6_4)   , layers_per_block=2   , in_channels=4   , down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   , cross_attention_dim=3_2   , conditioning_embedding_out_channels=(1_6, 3_2)   , )
    controlneta.controlnet_down_blocks.apply(a_     )

    torch.manual_seed(0     )
    UpperCAmelCase													=	ControlNetModel(
        block_out_channels=(3_2, 6_4)   , layers_per_block=2   , in_channels=4   , down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   , cross_attention_dim=3_2   , conditioning_embedding_out_channels=(1_6, 3_2)   , )
    controlneta.controlnet_down_blocks.apply(a_     )

    torch.manual_seed(0     )
    UpperCAmelCase													=	DDIMScheduler(
        beta_start=0.0_0085   , beta_end=0.012   , beta_schedule='scaled_linear'   , clip_sample=a_   , set_alpha_to_one=a_   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	AutoencoderKL(
        block_out_channels=[3_2, 6_4]   , in_channels=3   , out_channels=3   , down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   , up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   , latent_channels=4   , )
    torch.manual_seed(0     )
    UpperCAmelCase													=	CLIPTextConfig(
        bos_token_id=0   , eos_token_id=2   , hidden_size=3_2   , intermediate_size=3_7   , layer_norm_eps=1E-05   , num_attention_heads=4   , num_hidden_layers=5   , pad_token_id=1   , vocab_size=1_0_0_0   , )
    UpperCAmelCase													=	CLIPTextModel(a_     )
    UpperCAmelCase													=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

    UpperCAmelCase													=	MultiControlNetModel([controlneta, controlneta]     )

    UpperCAmelCase													=	{
        'unet': unet,
        'controlnet': controlnet,
        'scheduler': scheduler,
        'vae': vae,
        'text_encoder': text_encoder,
        'tokenizer': tokenizer,
        'safety_checker': None,
        'feature_extractor': None,
    }
    return components
  def 						snake_case_			(     self   , a_   , a_=0     )							->				int:


    """simple docstring"""
    if str(a_     ).startswith('mps'     ):
      UpperCAmelCase													=	torch.manual_seed(a_     )
    else:
      UpperCAmelCase													=	torch.Generator(device=a_     ).manual_seed(a_     )

    UpperCAmelCase													=	2

    UpperCAmelCase													=	[
        randn_tensor(
            (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)   , generator=a_   , device=torch.device(a_     )   , ),
        randn_tensor(
            (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)   , generator=a_   , device=torch.device(a_     )   , ),
    ]
    UpperCAmelCase													=	floats_tensor((1, 3, 3_2, 3_2)   , rng=random.Random(a_     )     ).to(a_     )
    UpperCAmelCase													=	init_image.cpu().permute(0   , 2   , 3   , 1     )[0]

    UpperCAmelCase													=	Image.fromarray(np.uinta(a_     )     ).convert('RGB'     ).resize((6_4, 6_4)     )
    UpperCAmelCase													=	Image.fromarray(np.uinta(init_image + 4     )     ).convert('RGB'     ).resize((6_4, 6_4)     )

    UpperCAmelCase													=	{
        'prompt': 'A painting of a squirrel eating a burger',
        'generator': generator,
        'num_inference_steps': 2,
        'guidance_scale': 6.0,
        'output_type': 'numpy',
        'image': image,
        'mask_image': mask_image,
        'control_image': control_image,
    }

    return inputs
  def 						snake_case_			(     self     )							->				Optional[Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.get_dummy_components()
    UpperCAmelCase													=	self.pipeline_class(**a_     )
    pipe.to(a_     )

    UpperCAmelCase													=	10.0
    UpperCAmelCase													=	4

    UpperCAmelCase													=	self.get_dummy_inputs(a_     )
    UpperCAmelCase													=	steps
    UpperCAmelCase													=	scale
    UpperCAmelCase													=	pipe(**a_     )[0]

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

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

    UpperCAmelCase													=	self.get_dummy_inputs(a_     )
    UpperCAmelCase													=	steps
    UpperCAmelCase													=	scale
    UpperCAmelCase													=	pipe(**a_   , 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 						snake_case_			(     self     )							->				Dict:


    """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 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3     )
  def 						snake_case_			(     self     )							->				Tuple:


    """simple docstring"""
    self._test_inference_batch_single_identical(expected_max_diff=2E-3     )





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


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




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

  '''simple docstring'''





  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()
  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    UpperCAmelCase													=	ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'     )

    UpperCAmelCase													=	StableDiffusionControlNetInpaintPipeline.from_pretrained(
        'runwayml/stable-diffusion-inpainting'   , safety_checker=a_   , controlnet=a_     )
    pipe.enable_model_cpu_offload()
    pipe.set_progress_bar_config(disable=a_     )

    UpperCAmelCase													=	torch.Generator(device='cpu'     ).manual_seed(0     )
    UpperCAmelCase													=	load_image(
        'https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'     ).resize((5_1_2, 5_1_2)     )

    UpperCAmelCase													=	load_image(
        'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
        '/stable_diffusion_inpaint/input_bench_mask.png'     ).resize((5_1_2, 5_1_2)     )

    UpperCAmelCase													=	'pitch black hole'

    UpperCAmelCase													=	load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'     ).resize((5_1_2, 5_1_2)     )

    UpperCAmelCase													=	pipe(
        a_   , image=a_   , mask_image=a_   , control_image=a_   , generator=a_   , output_type='np'   , num_inference_steps=3   , )

    UpperCAmelCase													=	output.images[0]

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

    UpperCAmelCase													=	load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/inpaint.npy'     )

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





  def 						snake_case_			(     self     )							->				int:


    """simple docstring"""
    UpperCAmelCase													=	ControlNetModel.from_pretrained('lllyasviel/control_v11p_sd15_inpaint'     )

    UpperCAmelCase													=	StableDiffusionControlNetInpaintPipeline.from_pretrained(
        'runwayml/stable-diffusion-v1-5'   , safety_checker=a_   , controlnet=a_     )
    UpperCAmelCase													=	DDIMScheduler.from_config(pipe.scheduler.config     )
    pipe.enable_model_cpu_offload()
    pipe.set_progress_bar_config(disable=a_     )

    UpperCAmelCase													=	torch.Generator(device='cpu'     ).manual_seed(3_3     )

    UpperCAmelCase													=	load_image(
        'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png'     )
    UpperCAmelCase													=	init_image.resize((5_1_2, 5_1_2)     )

    UpperCAmelCase													=	load_image(
        'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png'     )
    UpperCAmelCase													=	mask_image.resize((5_1_2, 5_1_2)     )

    UpperCAmelCase													=	'a handsome man with ray-ban sunglasses'

    def make_inpaint_condition(a_   , a_     ):
      UpperCAmelCase													=	np.array(image.convert('RGB'     )     ).astype(np.floataa     ) / 255.0
      UpperCAmelCase													=	np.array(image_mask.convert('L'     )     ).astype(np.floataa     ) / 255.0

      assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
      UpperCAmelCase													=	-1.0  # set as masked pixel
      UpperCAmelCase													=	np.expand_dims(a_   , 0     ).transpose(0   , 3   , 1   , 2     )
      UpperCAmelCase													=	torch.from_numpy(a_     )
      return image

    UpperCAmelCase													=	make_inpaint_condition(a_   , a_     )

    UpperCAmelCase													=	pipe(
        a_   , image=a_   , mask_image=a_   , control_image=a_   , guidance_scale=9.0   , eta=1.0   , generator=a_   , num_inference_steps=2_0   , output_type='np'   , )
    UpperCAmelCase													=	output.images[0]

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

    UpperCAmelCase													=	load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/boy_ray_ban.npy'     )

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





'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers.models.gptj.modeling_tf_gptj import (
      TFGPTJForCausalLM,
      TFGPTJForQuestionAnswering,
      TFGPTJForSequenceClassification,
      TFGPTJModel,
      shape_list,
  )



class        lowercase_	:

  '''simple docstring'''





  def __init__(     self   , a_     )							->				Any:


    """simple docstring"""
    UpperCAmelCase													=	parent
    UpperCAmelCase													=	1_3
    UpperCAmelCase													=	7
    UpperCAmelCase													=	True
    UpperCAmelCase													=	True
    UpperCAmelCase													=	True
    UpperCAmelCase													=	True
    UpperCAmelCase													=	True
    UpperCAmelCase													=	9_9
    UpperCAmelCase													=	3_2
    UpperCAmelCase													=	4
    UpperCAmelCase													=	2
    UpperCAmelCase													=	4
    UpperCAmelCase													=	3_7
    UpperCAmelCase													=	'gelu'
    UpperCAmelCase													=	0.1
    UpperCAmelCase													=	0.1
    UpperCAmelCase													=	5_1_2
    UpperCAmelCase													=	1_6
    UpperCAmelCase													=	2
    UpperCAmelCase													=	0.02
    UpperCAmelCase													=	3
    UpperCAmelCase													=	4
    UpperCAmelCase													=	None
    UpperCAmelCase													=	self.vocab_size - 1
    UpperCAmelCase													=	self.vocab_size - 1
    UpperCAmelCase													=	self.vocab_size - 1
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    UpperCAmelCase													=	ids_tensor([self.batch_size, self.seq_length]   , self.vocab_size     )

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

    UpperCAmelCase													=	None
    if self.use_token_type_ids:
      UpperCAmelCase													=	ids_tensor([self.batch_size, self.seq_length]   , self.type_vocab_size     )

    UpperCAmelCase													=	None
    if self.use_mc_token_ids:
      UpperCAmelCase													=	ids_tensor([self.batch_size, self.num_choices]   , self.seq_length     )

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

    UpperCAmelCase													=	GPTJConfig(
        vocab_size=self.vocab_size   , n_embd=self.hidden_size   , n_layer=self.num_hidden_layers   , n_head=self.num_attention_heads   , intermediate_size=self.intermediate_size   , hidden_act=self.hidden_act   , hidden_dropout_prob=self.hidden_dropout_prob   , attention_probs_dropout_prob=self.attention_probs_dropout_prob   , n_positions=self.max_position_embeddings   , type_vocab_size=self.type_vocab_size   , initializer_range=self.initializer_range   , bos_token_id=self.bos_token_id   , eos_token_id=self.eos_token_id   , pad_token_id=self.pad_token_id   , rotary_dim=self.rotary_dim   , return_dict=a_   , )

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

    return (
        config,
        input_ids,
        input_mask,
        head_mask,
        token_type_ids,
        mc_token_ids,
        sequence_labels,
        token_labels,
        choice_labels,
    )
  def 						snake_case_			(     self   , a_   , a_   , a_   , a_   , a_   , *a_     )							->				Dict:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJModel(config=a_     )
    UpperCAmelCase													=	{
        'input_ids': input_ids,
        'attention_mask': input_mask,
        'token_type_ids': token_type_ids,
    }
    UpperCAmelCase													=	model(a_     )

    UpperCAmelCase													=	[input_ids, None, input_mask]  # None is the input for 'past'
    UpperCAmelCase													=	model(a_     )

    UpperCAmelCase													=	model(a_     )

    self.parent.assertEqual(result.last_hidden_state.shape   , (self.batch_size, self.seq_length, self.hidden_size)     )
  def 						snake_case_			(     self   , a_   , a_   , a_   , a_   , a_   , *a_     )							->				Tuple:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJModel(config=a_     )

    # first forward pass
    UpperCAmelCase													=	model(a_   , token_type_ids=a_   , use_cache=a_     )
    UpperCAmelCase													=	model(a_   , token_type_ids=a_     )
    UpperCAmelCase													=	model(a_   , token_type_ids=a_   , use_cache=a_     )

    self.parent.assertTrue(len(a_     ) == len(a_     )     )
    self.parent.assertTrue(len(a_     ) == len(a_     ) + 1     )

    UpperCAmelCase							,      UpperCAmelCase													=	outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    UpperCAmelCase													=	ids_tensor((self.batch_size, 1)   , config.vocab_size     )
    UpperCAmelCase													=	ids_tensor([self.batch_size, 1]   , self.type_vocab_size     )

    # append to next input_ids and token_type_ids
    UpperCAmelCase													=	tf.concat([input_ids, next_tokens]   , axis=-1     )
    UpperCAmelCase													=	tf.concat([token_type_ids, next_token_types]   , axis=-1     )

    UpperCAmelCase													=	model(a_   , token_type_ids=a_     )['last_hidden_state']
    UpperCAmelCase													=	model(a_   , token_type_ids=a_   , past_key_values=a_     )[
        'last_hidden_state'
    ]

    # select random slice
    UpperCAmelCase													=	int(ids_tensor((1,)   , shape_list(a_     )[-1]     )     )
    UpperCAmelCase													=	output_from_no_past[:, -1, random_slice_idx]
    UpperCAmelCase													=	output_from_past[:, 0, random_slice_idx]

    # test that outputs are equal for slice
    tf.debugging.assert_near(a_   , a_   , rtol=1E-6     )
  def 						snake_case_			(     self   , a_   , a_   , a_   , a_   , a_   , *a_     )							->				Tuple:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJModel(config=a_     )

    # create attention mask
    UpperCAmelCase													=	self.seq_length // 2
    UpperCAmelCase													=	tf.ones((self.batch_size, half_seq_length)   , dtype=tf.intaa     )
    UpperCAmelCase													=	tf.zeros((self.batch_size, self.seq_length - half_seq_length)   , dtype=tf.intaa     )
    UpperCAmelCase													=	tf.concat([attn_mask_begin, attn_mask_end]   , axis=1     )

    # first forward pass
    UpperCAmelCase							,      UpperCAmelCase													=	model(a_   , attention_mask=a_     ).to_tuple()

    # create hypothetical next token and extent to next_input_ids
    UpperCAmelCase													=	ids_tensor((self.batch_size, 1)   , config.vocab_size     )

    # change a random masked slice from input_ids
    UpperCAmelCase													=	ids_tensor((1,)   , a_     ).numpy() + 1
    UpperCAmelCase													=	ids_tensor((self.batch_size, self.seq_length)   , config.vocab_size     )
    UpperCAmelCase													=	tf.range(self.seq_length     ) == (self.seq_length - random_seq_idx_to_change)
    UpperCAmelCase													=	tf.transpose(
        tf.broadcast_to(tf.expand_dims(a_   , -1     )   , (self.seq_length, self.batch_size)     )     )
    UpperCAmelCase													=	tf.where(a_   , a_   , a_     )

    # append to next input_ids and attn_mask
    UpperCAmelCase													=	tf.concat([input_ids, next_tokens]   , axis=-1     )
    UpperCAmelCase													=	tf.concat([attn_mask, tf.ones((shape_list(a_     )[0], 1)   , dtype=tf.intaa     )]   , axis=1     )

    # get two different outputs
    UpperCAmelCase													=	model(a_   , attention_mask=a_     )['last_hidden_state']
    UpperCAmelCase													=	model(a_   , past_key_values=a_   , attention_mask=a_     )[
        'last_hidden_state'
    ]

    # select random slice
    UpperCAmelCase													=	int(ids_tensor((1,)   , shape_list(a_     )[-1]     )     )
    UpperCAmelCase													=	output_from_no_past[:, -1, random_slice_idx]
    UpperCAmelCase													=	output_from_past[:, 0, random_slice_idx]

    # test that outputs are equal for slice
    tf.debugging.assert_near(a_   , a_   , rtol=1E-12     )
  def 						snake_case_			(     self   , a_   , a_   , a_   , a_   , a_   , *a_     )							->				Any:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJModel(config=a_     )

    UpperCAmelCase													=	input_ids[:1, :]
    UpperCAmelCase													=	input_mask[:1, :]
    UpperCAmelCase													=	token_type_ids[:1, :]
    UpperCAmelCase													=	1

    # first forward pass
    UpperCAmelCase													=	model(a_   , attention_mask=a_   , token_type_ids=a_   , use_cache=a_     )

    UpperCAmelCase							,      UpperCAmelCase													=	outputs.to_tuple()

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

    # append to next input_ids and token_type_ids
    UpperCAmelCase													=	tf.concat([input_ids, next_tokens]   , axis=-1     )
    UpperCAmelCase													=	tf.concat([input_mask, next_attn_mask]   , axis=-1     )
    UpperCAmelCase													=	tf.concat([token_type_ids, next_token_types]   , axis=-1     )

    UpperCAmelCase													=	model(
        a_   , token_type_ids=a_   , attention_mask=a_     )['last_hidden_state']
    UpperCAmelCase													=	model(
        a_   , token_type_ids=a_   , attention_mask=a_   , past_key_values=a_   , )['last_hidden_state']
    self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]     )

    # select random slice
    UpperCAmelCase													=	int(ids_tensor((1,)   , shape_list(a_     )[-1]     )     )
    UpperCAmelCase													=	output_from_no_past[:, -3:, random_slice_idx]
    UpperCAmelCase													=	output_from_past[:, :, random_slice_idx]

    # test that outputs are equal for slice
    tf.debugging.assert_near(a_   , a_   , rtol=1E-3     )
  def 						snake_case_			(     self   , a_   , a_   , a_   , a_   , a_   , *a_     )							->				Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJForCausalLM(config=a_     )
    UpperCAmelCase													=	{
        'input_ids': input_ids,
        'attention_mask': input_mask,
        'token_type_ids': token_type_ids,
    }
    UpperCAmelCase													=	model(a_     )
    self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.vocab_size)     )





  def 						snake_case_			(     self     )							->				Dict:


    """simple docstring"""
    UpperCAmelCase													=	self.prepare_config_and_inputs()

    (
        (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      (
        UpperCAmelCase
    )							,      
    )													=	config_and_inputs

    UpperCAmelCase													=	{
        'input_ids': input_ids,
        'token_type_ids': token_type_ids,
        'attention_mask': input_mask,
    }
    return config, inputs_dict




@require_tf
class        lowercase_	(		a ,       a ,       a ,       unittest.TestCase ):

  '''simple docstring'''





  __lowerCAmelCase					: Any      =   (
      (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
      if is_tf_available()
      else ()
  )

  __lowerCAmelCase					: Optional[int]      =   (TFGPTJForCausalLM,) if is_tf_available() else ()
  __lowerCAmelCase					: Union[str, Any]      =   (
      {
          "feature-extraction": TFGPTJModel,
          "question-answering": TFGPTJForQuestionAnswering,
          "text-classification": TFGPTJForSequenceClassification,
          "text-generation": TFGPTJForCausalLM,
          "zero-shot": TFGPTJForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  __lowerCAmelCase					: int      =   False
  __lowerCAmelCase					: List[Any]      =   False
  __lowerCAmelCase					: Any      =   False
  __lowerCAmelCase					: Tuple      =   False
  def 						snake_case_			(     self   , a_   , a_   , a_   , a_   , a_     )							->				int:


    """simple docstring"""
    if (
        pipeline_test_casse_name == "QAPipelineTests"
        and tokenizer_name is not None
        and not tokenizer_name.endswith('Fast'     )
    ):
      # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
      # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
      # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
      return True

    return False
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJModelTester(self     )
    UpperCAmelCase													=	ConfigTester(self   , config_class=a_   , n_embd=3_7     )
  def 						snake_case_			(     self     )							->				Any:


    """simple docstring"""
    self.config_tester.run_common_tests()
  def 						snake_case_			(     self     )							->				Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_gptj_model(*a_     )
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_gptj_model_past(*a_     )
  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_gptj_model_attention_mask_past(*a_     )
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_gptj_model_past_large_inputs(*a_     )
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_gptj_lm_head_model(*a_     )
  @slow
  @unittest.skipIf(
      not is_tf_available() or len(tf.config.list_physical_devices('GPU'     )     ) > 0   , 'skip testing on GPU for now to avoid GPU OOM.'   , )
  def 						snake_case_			(     self     )							->				int:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJModel.from_pretrained('EleutherAI/gpt-j-6B'   , from_pt=a_     )
    self.assertIsNotNone(a_     )





  @unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.'     )
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    super().test_resize_token_embeddings()




@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class        lowercase_	(		unittest.TestCase ):

  '''simple docstring'''





  def 						snake_case_			(     self     )							->				Any:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'   , from_pt=a_     )
    UpperCAmelCase													=	tf.convert_to_tensor([[4_6_4, 3_2_9_0]]   , dtype=tf.intaa     )  # The dog
    # fmt: off
    # The dog is a man's best friend. It is a loyal companion, and it is a friend
    UpperCAmelCase													=	[4_6_4, 3_2_9_0, 3_1_8, 2_5_7, 5_8_2, 3_3_8, 1_2_6_6, 1_5_4_5, 1_3, 6_3_2, 3_1_8, 2_5_7, 9_1_1_2, 1_5_1_8_5, 1_1, 2_9_0, 3_4_0, 3_1_8, 2_5_7, 1_5_4_5]
    # fmt: on
    UpperCAmelCase													=	model.generate(a_   , do_sample=a_     )
    self.assertListEqual(output_ids[0].numpy().tolist()   , a_     )
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    UpperCAmelCase													=	AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'   , revision='float16'     )
    UpperCAmelCase													=	TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'   , revision='float16'   , from_pt=a_     )

    UpperCAmelCase													=	tokenizer('Today is a nice day and'   , return_tensors='tf'     )
    # forces the generation to happen on CPU, to avoid GPU-related quirks
    with tf.device(':/CPU:0'     ):
      UpperCAmelCase													=	model.generate(**a_   , do_sample=a_   , seed=[4_2, 0]     )
    UpperCAmelCase													=	tokenizer.decode(output_ids[0]   , skip_special_tokens=a_     )

    UpperCAmelCase													=	'Today is a nice day and I’m going to go for a walk. I’'
    self.assertEqual(a_   , a_     )
  def 						snake_case_			(     self     )							->				Dict:


    """simple docstring"""
    UpperCAmelCase													=	TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'   , revision='float16'   , from_pt=a_     )
    UpperCAmelCase													=	AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'   , revision='float16'     )

    UpperCAmelCase													=	'left'

    # Define PAD Token = EOS Token = 50256
    UpperCAmelCase													=	tokenizer.eos_token
    UpperCAmelCase													=	model.config.eos_token_id

    # use different length sentences to test batching
    UpperCAmelCase													=	[
        'Hello, my dog is a little',
        'Today, I',
    ]
    UpperCAmelCase													=	[
        'Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia',
        'Today, I’m going to be talking about a topic that’',
    ]
    return model, tokenizer, sentences, expected_output_sentences
  def 						snake_case_			(     self     )							->				int:


    """simple docstring"""
    # Confirms that we get the expected results with left-padded beam search
    UpperCAmelCase							,      UpperCAmelCase							,      UpperCAmelCase							,      UpperCAmelCase													=	self._get_beam_search_test_objects()

    UpperCAmelCase													=	tokenizer(a_   , return_tensors='tf'   , padding=a_     )
    UpperCAmelCase													=	model.generate(**a_   , do_sample=a_   , num_beams=2     )
    UpperCAmelCase													=	tokenizer.batch_decode(a_   , skip_special_tokens=a_     )
    self.assertListEqual(a_   , a_     )
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    # Confirms that left-padding is working properly
    UpperCAmelCase							,      UpperCAmelCase							,      UpperCAmelCase							,      UpperCAmelCase													=	self._get_beam_search_test_objects()

    UpperCAmelCase													=	tokenizer(a_   , return_tensors='tf'   , padding=a_     )
    UpperCAmelCase													=	tokenizer(sentences[0]   , return_tensors='tf'     )
    UpperCAmelCase													=	model.generate(**a_   , do_sample=a_   , num_beams=2     )
    UpperCAmelCase													=	(
        shape_list(inputs_non_padded['input_ids']     )[-1]
        - tf.reduce_sum(tf.cast(inputs['attention_mask'][-1]   , tf.intaa     )     ).numpy()
    )
    UpperCAmelCase													=	tokenizer(sentences[1]   , return_tensors='tf'     )
    UpperCAmelCase													=	model.generate(
        **a_   , do_sample=a_   , num_beams=2   , max_length=model.config.max_length - num_paddings     )
    UpperCAmelCase													=	tokenizer.decode(output_non_padded[0]   , skip_special_tokens=a_     )
    UpperCAmelCase													=	tokenizer.decode(output_padded[0]   , skip_special_tokens=a_     )
    self.assertListEqual(a_   , [non_padded_sentence, padded_sentence]     )





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


    """simple docstring"""
    # Confirms that XLA is working properly
    UpperCAmelCase							,      UpperCAmelCase							,      UpperCAmelCase							,      UpperCAmelCase													=	self._get_beam_search_test_objects()

    UpperCAmelCase													=	tokenizer(a_   , return_tensors='tf'   , padding=a_     )
    UpperCAmelCase													=	tf.function(model.generate   , jit_compile=a_     )
    UpperCAmelCase													=	xla_generate(**a_   , do_sample=a_   , num_beams=2     )
    UpperCAmelCase													=	tokenizer.batch_decode(a_   , skip_special_tokens=a_     )
    self.assertListEqual(a_   , a_     )


