






import gc
import random
import unittest

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

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 							__UpperCamelCase     (  _lowerCAmelCase ,      _lowerCAmelCase ,      _lowerCAmelCase ,      unittest.TestCase ):
	__snake_case						:Union[str, Any] =							StableDiffusionInpaintPipeline
	__snake_case						:Optional[int] =							TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
	__snake_case						:Union[str, Any] =							TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
	__snake_case						:Any =							frozenset(
	    [] )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
	__snake_case						:Dict =							frozenset([] )







	def 			_a (	self   :				Any      )		->			Optional[Any]:
		"""simple docstring"""







		torch.manual_seed(0      )
		__lowercase   							=	UNetaDConditionModel(
		    block_out_channels=(32, 64)							,					layers_per_block=2							,					sample_size=32							,					in_channels=9							,					out_channels=4							,					down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")							,					up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")							,					cross_attention_dim=32							,					attention_head_dim=(2, 4)							,					use_linear_projection=_lowerCAmelCase							,					)
		__lowercase   							=	PNDMScheduler(skip_prk_steps=_lowerCAmelCase      )
		torch.manual_seed(0      )
		__lowercase   							=	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							,					sample_size=128							,					)
		torch.manual_seed(0      )
		__lowercase   							=	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=1000							,					hidden_act="""gelu"""							,					projection_dim=512							,					)
		__lowercase   							=	CLIPTextModel(_lowerCAmelCase      )
		__lowercase   							=	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

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







	def 			_a (	self   :				Tuple							,					_lowerCAmelCase   :				Optional[int]							,					_lowerCAmelCase   :				Optional[int]=0      )		->			Any:
		"""simple docstring"""







		__lowercase   							=	floats_tensor((1, 3, 32, 32)							,					rng=random.Random(_lowerCAmelCase      )      ).to(_lowerCAmelCase      )
		__lowercase   							=	image.cpu().permute(0							,					2							,					3							,					1      )[0]
		__lowercase   							=	Image.fromarray(np.uinta(_lowerCAmelCase      )      ).convert("""RGB"""      ).resize((64, 64)      )
		__lowercase   							=	Image.fromarray(np.uinta(image + 4      )      ).convert("""RGB"""      ).resize((64, 64)      )
		if str(_lowerCAmelCase      ).startswith("""mps"""      ):
			__lowercase   							=	torch.manual_seed(_lowerCAmelCase      )
		else:
			__lowercase   							=	torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
		__lowercase   							=	{
		    """prompt""": """A painting of a squirrel eating a burger""",
		    """image""": init_image,
		    """mask_image""": mask_image,
		    """generator""": generator,
		    """num_inference_steps""": 2,
		    """guidance_scale""": 6.0,
		    """output_type""": """numpy""",
		}
		return inputs







	def 			_a (	self   :				int      )		->			Union[str, Any]:
		"""simple docstring"""







		__lowercase   							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	StableDiffusionInpaintPipeline(**_lowerCAmelCase      )
		__lowercase   							=	sd_pipe.to(_lowerCAmelCase      )
		sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	sd_pipe(**_lowerCAmelCase      ).images
		__lowercase   							=	image[0, -3:, -3:, -1]

		assert image.shape == (1, 64, 64, 3)
		__lowercase   							=	np.array([0.4_727, 0.5_735, 0.3_941, 0.5_446, 0.5_926, 0.4_394, 0.5_062, 0.4_654, 0.4_476]      )

		assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2







	def 			_a (	self   :				Any      )		->			int:
		"""simple docstring"""







		super().test_inference_batch_single_identical(expected_max_diff=3e-3      )


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







	def 			_a (	self   :				Any      )		->			Any:
		"""simple docstring"""







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







	def 			_a (	self   :				List[Any]      )		->			Tuple:
		"""simple docstring"""







		__lowercase   							=	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/sd2-inpaint/init_image.png"""      )
		__lowercase   							=	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""      )
		__lowercase   							=	load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
		    """/yellow_cat_sitting_on_a_park_bench.npy"""      )

		__lowercase   							=	"""stabilityai/stable-diffusion-2-inpainting"""
		__lowercase   							=	StableDiffusionInpaintPipeline.from_pretrained(_lowerCAmelCase							,					safety_checker=_lowerCAmelCase      )
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing()

		__lowercase   							=	"""Face of a yellow cat, high resolution, sitting on a park bench"""

		__lowercase   							=	torch.manual_seed(0      )
		__lowercase   							=	pipe(
		    prompt=_lowerCAmelCase							,					image=_lowerCAmelCase							,					mask_image=_lowerCAmelCase							,					generator=_lowerCAmelCase							,					output_type="""np"""							,					)
		__lowercase   							=	output.images[0]

		assert image.shape == (512, 512, 3)
		assert np.abs(expected_image - image      ).max() < 9e-3







	def 			_a (	self   :				int      )		->			List[Any]:
		"""simple docstring"""







		__lowercase   							=	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/sd2-inpaint/init_image.png"""      )
		__lowercase   							=	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""      )
		__lowercase   							=	load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
		    """/yellow_cat_sitting_on_a_park_bench_fp16.npy"""      )

		__lowercase   							=	"""stabilityai/stable-diffusion-2-inpainting"""
		__lowercase   							=	StableDiffusionInpaintPipeline.from_pretrained(
		    _lowerCAmelCase							,					torch_dtype=torch.floataa							,					safety_checker=_lowerCAmelCase							,					)
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing()

		__lowercase   							=	"""Face of a yellow cat, high resolution, sitting on a park bench"""

		__lowercase   							=	torch.manual_seed(0      )
		__lowercase   							=	pipe(
		    prompt=_lowerCAmelCase							,					image=_lowerCAmelCase							,					mask_image=_lowerCAmelCase							,					generator=_lowerCAmelCase							,					output_type="""np"""							,					)
		__lowercase   							=	output.images[0]

		assert image.shape == (512, 512, 3)
		assert np.abs(expected_image - image      ).max() < 5e-1







	def 			_a (	self   :				Optional[Any]      )		->			Union[str, Any]:
		"""simple docstring"""







		torch.cuda.empty_cache()
		torch.cuda.reset_max_memory_allocated()
		torch.cuda.reset_peak_memory_stats()

		__lowercase   							=	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/sd2-inpaint/init_image.png"""      )
		__lowercase   							=	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""      )

		__lowercase   							=	"""stabilityai/stable-diffusion-2-inpainting"""
		__lowercase   							=	PNDMScheduler.from_pretrained(_lowerCAmelCase							,					subfolder="""scheduler"""      )
		__lowercase   							=	StableDiffusionInpaintPipeline.from_pretrained(
		    _lowerCAmelCase							,					safety_checker=_lowerCAmelCase							,					scheduler=_lowerCAmelCase							,					torch_dtype=torch.floataa							,					)
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing(1      )
		pipe.enable_sequential_cpu_offload()

		__lowercase   							=	"""Face of a yellow cat, high resolution, sitting on a park bench"""

		__lowercase   							=	torch.manual_seed(0      )
		__lowercase   							=	pipe(
		    prompt=_lowerCAmelCase							,					image=_lowerCAmelCase							,					mask_image=_lowerCAmelCase							,					generator=_lowerCAmelCase							,					num_inference_steps=2							,					output_type="""np"""							,					)

		__lowercase   							=	torch.cuda.max_memory_allocated()
		# make sure that less than 2.65 GB is allocated
		assert mem_bytes < 2.65 * 10**9





"""simple docstring"""


import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 (
	    FalconForCausalLM,
	    FalconForQuestionAnswering,
	    FalconForSequenceClassification,
	    FalconForTokenClassification,
	    FalconModel,
	)
class     _UpperCamelCase    :
							'''simple docstring'''







							def __init__(	self       ,  __a       ,  __a=3       ,  __a=7       ,  __a=True       ,  __a=True       ,  __a=False       ,  __a=True       ,  __a=99       ,  __a=32       ,  __a=5       ,  __a=4       ,  __a=37       ,  __a="gelu"       ,  __a=0.1       ,  __a=0.1       ,  __a=5_12       ,  __a=16       ,  __a=2       ,  __a=0.0_2       ,  __a=3       ,  __a=4       ,  __a=None       ,  ):
										__lowerCAmelCase							       =     parent
										__lowerCAmelCase							       =     batch_size
										__lowerCAmelCase							       =     seq_length
										__lowerCAmelCase							       =     is_training
										__lowerCAmelCase							       =     use_input_mask
										__lowerCAmelCase							       =     use_token_type_ids
										__lowerCAmelCase							       =     use_labels
										__lowerCAmelCase							       =     vocab_size
										__lowerCAmelCase							       =     hidden_size
										__lowerCAmelCase							       =     num_hidden_layers
										__lowerCAmelCase							       =     num_attention_heads
										__lowerCAmelCase							       =     intermediate_size
										__lowerCAmelCase							       =     hidden_act
										__lowerCAmelCase							       =     hidden_dropout_prob
										__lowerCAmelCase							       =     attention_probs_dropout_prob
										__lowerCAmelCase							       =     max_position_embeddings
										__lowerCAmelCase							       =     type_vocab_size
										__lowerCAmelCase							       =     type_sequence_label_size
										__lowerCAmelCase							       =     initializer_range
										__lowerCAmelCase							       =     num_labels
										__lowerCAmelCase							       =     num_choices
										__lowerCAmelCase							       =     scope

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     ids_tensor([self.batch_size, self.seq_length]       ,  self.vocab_size      )

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

										__lowerCAmelCase							       =     None

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

										__lowerCAmelCase							       =     self.get_config()

										return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

							def      snake_case							(	self      ):
										return FalconConfig(
										    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_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       ,  max_position_embeddings=self.max_position_embeddings       ,  type_vocab_size=self.type_vocab_size       ,  is_decoder=__a       ,  initializer_range=self.initializer_range       ,  pad_token_id=1       ,  new_decoder_architecture=__a       ,  )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     FalconModel(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a      )
										__lowerCAmelCase							       =     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       ,  __a       ,  __a       ,  __a       ,  ):
										__lowerCAmelCase							       =     True
										__lowerCAmelCase							       =     FalconModel(__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  encoder_hidden_states=__a       ,  encoder_attention_mask=__a       ,  )
										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  encoder_hidden_states=__a       ,  )
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__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       ,  __a       ,  __a       ,  __a       ,  ):
										__lowerCAmelCase							       =     FalconForCausalLM(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  labels=__a      )
										self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)      )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  ):
										__lowerCAmelCase							       =     True
										__lowerCAmelCase							       =     True
										__lowerCAmelCase							       =     FalconForCausalLM(config=__a      )
										model.to(__a      )
										model.eval()

										# first forward pass
										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  encoder_hidden_states=__a       ,  encoder_attention_mask=__a       ,  use_cache=__a       ,  )
										__lowerCAmelCase							       =     outputs.past_key_values

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

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

										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  encoder_hidden_states=__a       ,  encoder_attention_mask=__a       ,  output_hidden_states=__a       ,  )["hidden_states"][0]
										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  encoder_hidden_states=__a       ,  encoder_attention_mask=__a       ,  past_key_values=__a       ,  output_hidden_states=__a       ,  )["hidden_states"][0]

										# select random slice
										__lowerCAmelCase							       =     ids_tensor((1,)       ,  output_from_past.shape[-1]      ).item()
										__lowerCAmelCase							       =     output_from_no_past[:, -3:, random_slice_idx].detach()
										__lowerCAmelCase							       =     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(__a       ,  __a       ,  atol=1e-3      )      )



							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.prepare_config_and_inputs()
										(
										    (
										    __lowerCAmelCase
										)       ,      (
										    __lowerCAmelCase
										)       ,      (
										    __lowerCAmelCase
										)       ,      (
										    __lowerCAmelCase
										)       ,      (
										    __lowerCAmelCase
										)       ,      (
										    __lowerCAmelCase
										)       ,      (
										    __lowerCAmelCase
										)       ,      
										)							       =     config_and_inputs
										__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask}
										return config, inputs_dict
@require_torch
class     _UpperCamelCase    (  lowerCAmelCase__			,lowerCAmelCase__			,lowerCAmelCase__			,unittest.TestCase						):
							'''simple docstring'''







							__UpperCAmelCase		:       int							 =(
							    (
							        FalconModel,
							        FalconForCausalLM,
							        FalconForSequenceClassification,
							        FalconForTokenClassification,
							        FalconForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							__UpperCAmelCase		:       Union[str, Any]							 =(FalconForCausalLM,) if is_torch_available() else ()
							__UpperCAmelCase		:       List[Any]							 =(
							    {
							        """feature-extraction""": FalconModel,
							        """text-classification""": FalconForSequenceClassification,
							        """text-generation""": FalconForCausalLM,
							        """question-answering""": FalconForQuestionAnswering,
							        """token-classification""": FalconForTokenClassification,
							        """zero-shot""": FalconForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							__UpperCAmelCase		:       Any							 =False
							__UpperCAmelCase		:       Tuple							 =False

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     FalconModelTester(self      )
										__lowerCAmelCase							       =     ConfigTester(self       ,  config_class=__a       ,  hidden_size=37      )

							def      snake_case							(	self      ):
										self.config_tester.run_common_tests()

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase       ,      *__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										for alibi in [True, False]:
													__lowerCAmelCase							       =     alibi
													self.model_tester.create_and_check_model(__a       ,  *__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase       ,      __lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase							       =     3
										__lowerCAmelCase							       =     input_dict["input_ids"]
										__lowerCAmelCase							       =     input_ids.ne(1      ).to(__a      )
										__lowerCAmelCase							       =     ids_tensor([self.model_tester.batch_size]       ,  self.model_tester.type_sequence_label_size      )
										__lowerCAmelCase							       =     FalconForSequenceClassification(__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  labels=__a      )
										self.assertEqual(result.logits.shape       ,  (self.model_tester.batch_size, self.model_tester.num_labels)      )

							def      snake_case							(	self      ):
										__lowerCAmelCase       ,      __lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase							       =     3
										__lowerCAmelCase							       =     "single_label_classification"
										__lowerCAmelCase							       =     input_dict["input_ids"]
										__lowerCAmelCase							       =     input_ids.ne(1      ).to(__a      )
										__lowerCAmelCase							       =     ids_tensor([self.model_tester.batch_size]       ,  self.model_tester.type_sequence_label_size      )
										__lowerCAmelCase							       =     FalconForSequenceClassification(__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  labels=__a      )
										self.assertEqual(result.logits.shape       ,  (self.model_tester.batch_size, self.model_tester.num_labels)      )

							def      snake_case							(	self      ):
										__lowerCAmelCase       ,      __lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase							       =     input_dict["input_ids"]
										__lowerCAmelCase							       =     FalconForCausalLM(__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  use_cache=__a      )
										__lowerCAmelCase							       =     input_ids.shape[0]
										__lowerCAmelCase							       =     model._convert_to_rw_cache(result.past_key_values      )
										__lowerCAmelCase							       =     model._convert_cache_to_standard_format(__a       ,  __a      )
										for layer in range(len(__a      )      ):
													for tensor_idx in range(2      ):
																self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3      )
																self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4      )
																self.assertTrue(
																    torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]      )      )

							def      snake_case							(	self      ):
										__lowerCAmelCase       ,      __lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase							       =     3
										__lowerCAmelCase							       =     "multi_label_classification"
										__lowerCAmelCase							       =     input_dict["input_ids"]
										__lowerCAmelCase							       =     input_ids.ne(1      ).to(__a      )
										__lowerCAmelCase							       =     ids_tensor(
										    [self.model_tester.batch_size, config.num_labels]       ,  self.model_tester.type_sequence_label_size      ).to(torch.float      )
										__lowerCAmelCase							       =     FalconForSequenceClassification(__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  labels=__a      )
										self.assertEqual(result.logits.shape       ,  (self.model_tester.batch_size, self.model_tester.num_labels)      )



							def      snake_case							(	self      ):
										# Falcon can have different numbers of KV-heads than the number of query heads, so we need
										# to override this test to use the right head counts.
										for model_class in self.all_generative_model_classes:
													__lowerCAmelCase       ,      __lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs_for_common()

													# If it doesn't support cache, pass the test
													if not hasattr(__a       ,  "use_cache"      ):
																return

													__lowerCAmelCase							       =     model_class(__a      ).to(__a      )
													if "use_cache" not in inputs:
																__lowerCAmelCase							       =     True
													__lowerCAmelCase							       =     model(**__a      )

													# If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
													if "past_key_values" not in outputs:
																return

													__lowerCAmelCase							       =     (
													    getattr(__a       ,  "decoder_layers"       ,  __a      )
													    or getattr(__a       ,  "num_decoder_layers"       ,  __a      )
													    or config.num_hidden_layers
													)
													__lowerCAmelCase							       =     getattr(__a       ,  "num_kv_heads"       ,  config.num_attention_heads      )
													__lowerCAmelCase							       =     getattr(__a       ,  "d_model"       ,  config.hidden_size      )
													__lowerCAmelCase							       =     embed_dim // num_attention_heads

													__lowerCAmelCase							       =     outputs["past_key_values"]
													self.assertEqual(len(__a      )       ,  __a      )

													__lowerCAmelCase       ,      __lowerCAmelCase							       =     inputs["input_ids"].shape
													for i in range(__a      ):
																if config.new_decoder_architecture:
																			__lowerCAmelCase							       =     config.num_attention_heads
																elif config.multi_query:
																			__lowerCAmelCase							       =     1
																self.assertEqual(len(past_kv[0]      )       ,  2      )  # K V for the decoder = 2
																self.assertEqual(
																    past_kv[i][0].shape       ,  (batch_size, num_attention_heads, seq_length, per_head_embed_dim)      )
																self.assertEqual(
																    past_kv[i][1].shape       ,  (batch_size, num_attention_heads, seq_length, per_head_embed_dim)      )
@require_torch
class     _UpperCamelCase    (  unittest.TestCase						):
							'''simple docstring'''







							@slow
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b"      )
										__lowerCAmelCase							       =     FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b"      )
										model.eval()
										model.to(__a      )
										__lowerCAmelCase							       =     tokenizer("My favorite food is"       ,  return_tensors="pt"      ).to(__a      )

										__lowerCAmelCase							       =     (
										    "My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."
										)

										__lowerCAmelCase							       =     model.generate(**__a       ,  do_sample=__a       ,  max_new_tokens=19      )
										__lowerCAmelCase							       =     tokenizer.batch_decode(__a      )[0]

										self.assertEqual(__a       ,  __a      )

							@slow
							def      snake_case							(	self      ):
										# The big models are way too big for the CI, so we use tiny random models that resemble their
										# architectures but with much smaller and fewer layers
										for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
													__lowerCAmelCase							       =     AutoTokenizer.from_pretrained(__a      )
													__lowerCAmelCase							       =     FalconForCausalLM.from_pretrained(__a      )
													model.eval()
													model.to(__a      )
													__lowerCAmelCase							       =     tokenizer("My favorite food is"       ,  return_tensors="pt"      ).to(__a      )

													# We just test that these run without errors - the models are randomly initialized
													# and so the actual text outputs will be garbage
													model.generate(**__a       ,  do_sample=__a       ,  max_new_tokens=4      )
													model.generate(**__a       ,  do_sample=__a       ,  max_new_tokens=4      )
													model.generate(**__a       ,  num_beams=2       ,  max_new_tokens=4      )



							@slow
							def      snake_case							(	self      ):
										# The big models are way too big for the CI, so we use tiny random models that resemble their
										# architectures but with much smaller and fewer layers
										with torch.no_grad():
													for repo in [
													    "Rocketknight1/falcon-rw-1b",
													    "Rocketknight1/tiny-random-falcon-7b",
													    "Rocketknight1/tiny-random-falcon-40b",
													]:
																__lowerCAmelCase							       =     AutoTokenizer.from_pretrained(__a      )
																__lowerCAmelCase							       =     FalconForCausalLM.from_pretrained(__a      )
																model.eval()
																model.to(device=__a      )
																__lowerCAmelCase							       =     tokenizer("My favorite food is"       ,  return_tensors="pt"      ).to(__a      )

																# Test results are the same with and without cache
																__lowerCAmelCase							       =     model.generate(**__a       ,  do_sample=__a       ,  max_new_tokens=20       ,  use_cache=__a      )
																__lowerCAmelCase							       =     model.generate(**__a       ,  do_sample=__a       ,  max_new_tokens=20       ,  use_cache=__a      )
																self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0      )




