



import argparse
import json
from pathlib import Path

import requests
import torch
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BeitConfig,
    BeitForImageClassification,
    BeitForMaskedImageModeling,
    BeitForSemanticSegmentation,
    BeitImageProcessor,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
a			:     Any     = logging.get_logger(__name__)







def    lowerCAmelCase_  (lowerCAmelCase__:		Union[str, Any]		,	lowerCAmelCase__:		Tuple=False		,	lowerCAmelCase__:		int=False     ):





			"""simple docstring"""


			UpperCAmelCase_:		List[str]       			=    """backbone.""" if is_semantic else """"""

			UpperCAmelCase_:		List[str]       			=    []
			for i in range(config.num_hidden_layers     ):
						# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
						rename_keys.append((F'{prefix}blocks.{i}.norm1.weight', F'beit.encoder.layer.{i}.layernorm_before.weight')     )
						rename_keys.append((F'{prefix}blocks.{i}.norm1.bias', F'beit.encoder.layer.{i}.layernorm_before.bias')     )
						rename_keys.append(
						    (F'{prefix}blocks.{i}.attn.proj.weight', F'beit.encoder.layer.{i}.attention.output.dense.weight')     )
						rename_keys.append(
						    (F'{prefix}blocks.{i}.attn.proj.bias', F'beit.encoder.layer.{i}.attention.output.dense.bias')     )
						rename_keys.append((F'{prefix}blocks.{i}.norm2.weight', F'beit.encoder.layer.{i}.layernorm_after.weight')     )
						rename_keys.append((F'{prefix}blocks.{i}.norm2.bias', F'beit.encoder.layer.{i}.layernorm_after.bias')     )
						rename_keys.append((F'{prefix}blocks.{i}.mlp.fc1.weight', F'beit.encoder.layer.{i}.intermediate.dense.weight')     )
						rename_keys.append((F'{prefix}blocks.{i}.mlp.fc1.bias', F'beit.encoder.layer.{i}.intermediate.dense.bias')     )
						rename_keys.append((F'{prefix}blocks.{i}.mlp.fc2.weight', F'beit.encoder.layer.{i}.output.dense.weight')     )
						rename_keys.append((F'{prefix}blocks.{i}.mlp.fc2.bias', F'beit.encoder.layer.{i}.output.dense.bias')     )

			# projection layer + position embeddings
			rename_keys.extend(
			    [
			        (F'{prefix}cls_token', """beit.embeddings.cls_token"""),
			        (F'{prefix}patch_embed.proj.weight', """beit.embeddings.patch_embeddings.projection.weight"""),
			        (F'{prefix}patch_embed.proj.bias', """beit.embeddings.patch_embeddings.projection.bias"""),
			    ]     )

			if has_lm_head:
						# mask token + shared relative position bias + layernorm
						rename_keys.extend(
						    [
						        ("""mask_token""", """beit.embeddings.mask_token"""),
						        (
						            """rel_pos_bias.relative_position_bias_table""",
						            """beit.encoder.relative_position_bias.relative_position_bias_table""",
						        ),
						        (
						            """rel_pos_bias.relative_position_index""",
						            """beit.encoder.relative_position_bias.relative_position_index""",
						        ),
						        ("""norm.weight""", """layernorm.weight"""),
						        ("""norm.bias""", """layernorm.bias"""),
						    ]     )
			elif is_semantic:
						# semantic segmentation classification heads
						rename_keys.extend(
						    [
						        ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
						        ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
						        ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
						        ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
						    ]     )
			else:
						# layernorm + classification head
						rename_keys.extend(
						    [
						        ("""fc_norm.weight""", """beit.pooler.layernorm.weight"""),
						        ("""fc_norm.bias""", """beit.pooler.layernorm.bias"""),
						        ("""head.weight""", """classifier.weight"""),
						        ("""head.bias""", """classifier.bias"""),
						    ]     )

			return rename_keys







def    lowerCAmelCase_  (lowerCAmelCase__:		Any		,	lowerCAmelCase__:		str		,	lowerCAmelCase__:		Optional[Any]=False		,	lowerCAmelCase__:		Union[str, Any]=False     ):





			"""simple docstring"""


			for i in range(config.num_hidden_layers     ):
						UpperCAmelCase_:		str       			=    """backbone.""" if is_semantic else """"""
						# queries, keys and values
						UpperCAmelCase_:		List[str]       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.qkv.weight'     )
						UpperCAmelCase_:		Optional[Any]       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.q_bias'     )
						UpperCAmelCase_:		Optional[int]       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.v_bias'     )

						UpperCAmelCase_:		List[Any]       			=    in_proj_weight[
						    : config.hidden_size, :
						]
						UpperCAmelCase_:		List[str]       			=    q_bias
						UpperCAmelCase_:		int       			=    in_proj_weight[
						    config.hidden_size : config.hidden_size * 2, :
						]
						UpperCAmelCase_:		Tuple       			=    in_proj_weight[
						    -config.hidden_size :, :
						]
						UpperCAmelCase_:		Tuple       			=    v_bias

						# gamma_1 and gamma_2
						# we call them lambda because otherwise they are renamed when using .from_pretrained
						UpperCAmelCase_:		Any       			=    state_dict.pop(F'{prefix}blocks.{i}.gamma_1'     )
						UpperCAmelCase_:		List[Any]       			=    state_dict.pop(F'{prefix}blocks.{i}.gamma_2'     )

						UpperCAmelCase_:		Dict       			=    gamma_a
						UpperCAmelCase_:		str       			=    gamma_a

						# relative_position bias table + index
						if not has_lm_head:
									# each layer has its own relative position bias
									UpperCAmelCase_:		Dict       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.relative_position_bias_table'     )
									UpperCAmelCase_:		List[Any]       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.relative_position_index'     )

									UpperCAmelCase_:		List[str]       			=    table
									UpperCAmelCase_:		List[str]       			=    index







def    lowerCAmelCase_  (lowerCAmelCase__:		Optional[int]		,	lowerCAmelCase__:		Optional[Any]		,	lowerCAmelCase__:		Union[str, Any]     ):





			"""simple docstring"""


			UpperCAmelCase_:		List[str]       			=    dct.pop(lowerCAmelCase__     )
			UpperCAmelCase_:		Dict       			=    val







def    lowerCAmelCase_  ():





			"""simple docstring"""


			UpperCAmelCase_:		str       			=    """http://images.cocodataset.org/val2017/000000039769.jpg"""
			UpperCAmelCase_:		Dict       			=    Image.open(requests.get(lowerCAmelCase__		,	stream=lowerCAmelCase__     ).raw     )
			return im







@torch.no_grad()
def    lowerCAmelCase_  (lowerCAmelCase__:		List[Any]		,	lowerCAmelCase__:		List[str]     ):





			"""simple docstring"""


			UpperCAmelCase_:		Optional[int]       			=    BeitConfig()
			UpperCAmelCase_:		Optional[Any]       			=    False
			UpperCAmelCase_:		str       			=    False
			UpperCAmelCase_:		Union[str, Any]       			=    """huggingface/label-files"""
			# set config parameters based on URL
			if checkpoint_url[-9:-4] == "pt22k":
						# masked image modeling
						UpperCAmelCase_:		List[Any]       			=    True
						UpperCAmelCase_:		Union[str, Any]       			=    True
						UpperCAmelCase_:		List[str]       			=    True
			elif checkpoint_url[-9:-4] == "ft22k":
						# intermediate fine-tuning on ImageNet-22k
						UpperCAmelCase_:		str       			=    True
						UpperCAmelCase_:		Optional[int]       			=    2_1_8_4_1
						UpperCAmelCase_:		Dict       			=    """imagenet-22k-id2label.json"""
						UpperCAmelCase_:		Tuple       			=    json.load(open(hf_hub_download(lowerCAmelCase__		,	lowerCAmelCase__		,	repo_type="""dataset"""     )		,	"""r"""     )     )
						UpperCAmelCase_:		Optional[Any]       			=    {int(lowerCAmelCase__     ): v for k, v in idalabel.items()}
						# this dataset contains 21843 labels but the model only has 21841
						# we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
						del idalabel[9_2_0_5]
						del idalabel[1_5_0_2_7]
						UpperCAmelCase_:		Optional[Any]       			=    idalabel
						UpperCAmelCase_:		Optional[Any]       			=    {v: k for k, v in idalabel.items()}
			elif checkpoint_url[-8:-4] == "to1k":
						# fine-tuning on ImageNet-1k
						UpperCAmelCase_:		int       			=    True
						UpperCAmelCase_:		Any       			=    1_0_0_0
						UpperCAmelCase_:		List[str]       			=    """imagenet-1k-id2label.json"""
						UpperCAmelCase_:		Dict       			=    json.load(open(hf_hub_download(lowerCAmelCase__		,	lowerCAmelCase__		,	repo_type="""dataset"""     )		,	"""r"""     )     )
						UpperCAmelCase_:		List[Any]       			=    {int(lowerCAmelCase__     ): v for k, v in idalabel.items()}
						UpperCAmelCase_:		List[Any]       			=    idalabel
						UpperCAmelCase_:		int       			=    {v: k for k, v in idalabel.items()}
						if "384" in checkpoint_url:
									UpperCAmelCase_:		Union[str, Any]       			=    3_8_4
						if "512" in checkpoint_url:
									UpperCAmelCase_:		Optional[Any]       			=    5_1_2
			elif "ade20k" in checkpoint_url:
						# fine-tuning
						UpperCAmelCase_:		Any       			=    True
						UpperCAmelCase_:		List[str]       			=    1_5_0
						UpperCAmelCase_:		Dict       			=    """ade20k-id2label.json"""
						UpperCAmelCase_:		List[str]       			=    json.load(open(hf_hub_download(lowerCAmelCase__		,	lowerCAmelCase__		,	repo_type="""dataset"""     )		,	"""r"""     )     )
						UpperCAmelCase_:		Tuple       			=    {int(lowerCAmelCase__     ): v for k, v in idalabel.items()}
						UpperCAmelCase_:		Tuple       			=    idalabel
						UpperCAmelCase_:		Tuple       			=    {v: k for k, v in idalabel.items()}
						UpperCAmelCase_:		List[str]       			=    6_4_0
						UpperCAmelCase_:		List[str]       			=    True
			else:
						raise ValueError("""Checkpoint not supported, URL should either end with 'pt22k', 'ft22k', 'to1k' or 'ade20k'"""     )

			# size of the architecture
			if "base" in checkpoint_url:
						pass
			elif "large" in checkpoint_url:
						UpperCAmelCase_:		str       			=    1_0_2_4
						UpperCAmelCase_:		Optional[Any]       			=    4_0_9_6
						UpperCAmelCase_:		Optional[Any]       			=    2_4
						UpperCAmelCase_:		str       			=    1_6
						if "ade20k" in checkpoint_url:
									UpperCAmelCase_:		Dict       			=    6_4_0
									UpperCAmelCase_:		Union[str, Any]       			=    [7, 1_1, 1_5, 2_3]
			else:
						raise ValueError("""Should either find 'base' or 'large' in checkpoint URL"""     )

			# load state_dict of original model, remove and rename some keys
			UpperCAmelCase_:		Optional[Any]       			=    torch.hub.load_state_dict_from_url(lowerCAmelCase__		,	map_location="""cpu"""		,	check_hash=lowerCAmelCase__     )
			UpperCAmelCase_:		List[str]       			=    state_dict["""model"""] if """ade20k""" not in checkpoint_url else state_dict["""state_dict"""]

			UpperCAmelCase_:		List[str]       			=    create_rename_keys(lowerCAmelCase__		,	has_lm_head=lowerCAmelCase__		,	is_semantic=lowerCAmelCase__     )
			for src, dest in rename_keys:
						rename_key(lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__     )
			read_in_q_k_v(lowerCAmelCase__		,	lowerCAmelCase__		,	has_lm_head=lowerCAmelCase__		,	is_semantic=lowerCAmelCase__     )
			if is_semantic:
						# add prefix to decoder keys
						for key, val in state_dict.copy().items():
									UpperCAmelCase_:		Any       			=    state_dict.pop(lowerCAmelCase__     )
									if key.startswith("""backbone.fpn"""     ):
												UpperCAmelCase_:		Union[str, Any]       			=    key.replace("""backbone.fpn"""		,	"""fpn"""     )
									UpperCAmelCase_:		Dict       			=    val

    # load HuggingFace model
			if checkpoint_url[-9:-4] == "pt22k":
						UpperCAmelCase_:		List[str]       			=    BeitForMaskedImageModeling(lowerCAmelCase__     )
			elif "ade20k" in checkpoint_url:
						UpperCAmelCase_:		Dict       			=    BeitForSemanticSegmentation(lowerCAmelCase__     )
			else:
						UpperCAmelCase_:		List[Any]       			=    BeitForImageClassification(lowerCAmelCase__     )
			model.eval()
			model.load_state_dict(lowerCAmelCase__     )

			# Check outputs on an image
			if is_semantic:
						UpperCAmelCase_:		Optional[Any]       			=    BeitImageProcessor(size=config.image_size		,	do_center_crop=lowerCAmelCase__     )
						UpperCAmelCase_:		List[str]       			=    load_dataset("""hf-internal-testing/fixtures_ade20k"""		,	split="""test"""     )
						UpperCAmelCase_:		Union[str, Any]       			=    Image.open(ds[0]["""file"""]     )
			else:
						UpperCAmelCase_:		List[Any]       			=    BeitImageProcessor(
						    size=config.image_size		,	resample=PILImageResampling.BILINEAR		,	do_center_crop=lowerCAmelCase__     )
						UpperCAmelCase_:		Any       			=    prepare_img()

			UpperCAmelCase_:		Tuple       			=    image_processor(images=lowerCAmelCase__		,	return_tensors="""pt"""     )
			UpperCAmelCase_:		Optional[Any]       			=    encoding["""pixel_values"""]

			UpperCAmelCase_:		Dict       			=    model(lowerCAmelCase__     )
			UpperCAmelCase_:		Optional[Any]       			=    outputs.logits

			# verify logits
			UpperCAmelCase_:		List[str]       			=    torch.Size([1, 1_0_0_0]     )
			if checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k"""     ):
						UpperCAmelCase_:		Union[str, Any]       			=    torch.Size([1, 1_9_6, 8_1_9_2]     )
			elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k"""     ):
						UpperCAmelCase_:		Union[str, Any]       			=    torch.Size([1, 1_9_6, 8_1_9_2]     )
			elif checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k_ft22k"""     ):
						UpperCAmelCase_:		Optional[Any]       			=    torch.Size([1, 2_1_8_4_1]     )
						UpperCAmelCase_:		List[Any]       			=    torch.tensor([2.2288, 2.4671, 0.7395]     )
						UpperCAmelCase_:		List[str]       			=    2_3_9_7
			elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k_ft22k"""     ):
						UpperCAmelCase_:		str       			=    torch.Size([1, 2_1_8_4_1]     )
						UpperCAmelCase_:		Optional[Any]       			=    torch.tensor([1.6881, -0.2787, 0.5901]     )
						UpperCAmelCase_:		Optional[Any]       			=    2_3_9_6
			elif checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k_ft1k"""     ):
						UpperCAmelCase_:		List[str]       			=    torch.tensor([0.1241, 0.0798, -0.6569]     )
						UpperCAmelCase_:		int       			=    2_8_5
			elif checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k_ft22kto1k"""     ):
						UpperCAmelCase_:		List[str]       			=    torch.tensor([-1.2385, -1.0987, -1.0108]     )
						UpperCAmelCase_:		int       			=    2_8_1
			elif checkpoint_url[:-4].endswith("""beit_base_patch16_384_pt22k_ft22kto1k"""     ):
						UpperCAmelCase_:		Optional[int]       			=    torch.tensor([-1.5303, -0.9484, -0.3147]     )
						UpperCAmelCase_:		Optional[int]       			=    7_6_1
			elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k_ft1k"""     ):
						UpperCAmelCase_:		Union[str, Any]       			=    torch.tensor([0.4610, -0.0928, 0.2086]     )
						UpperCAmelCase_:		int       			=    7_6_1
			elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k_ft22kto1k"""     ):
						UpperCAmelCase_:		int       			=    torch.tensor([-0.4804, 0.6257, -0.1837]     )
						UpperCAmelCase_:		int       			=    7_6_1
			elif checkpoint_url[:-4].endswith("""beit_large_patch16_384_pt22k_ft22kto1k"""     ):
						UpperCAmelCase_:		Any       			=    torch.tensor([[-0.5122, 0.5117, -0.2113]]     )
						UpperCAmelCase_:		str       			=    7_6_1
			elif checkpoint_url[:-4].endswith("""beit_large_patch16_512_pt22k_ft22kto1k"""     ):
						UpperCAmelCase_:		int       			=    torch.tensor([-0.3062, 0.7261, 0.4852]     )
						UpperCAmelCase_:		Tuple       			=    7_6_1
			elif checkpoint_url[:-4].endswith("""beit_base_patch16_640_pt22k_ft22ktoade20k"""     ):
						UpperCAmelCase_:		Dict       			=    (1, 1_5_0, 1_6_0, 1_6_0)
						UpperCAmelCase_:		Union[str, Any]       			=    torch.tensor(
						    [
						        [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
						        [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
						        [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
						    ]     )
			elif checkpoint_url[:-4].endswith("""beit_large_patch16_640_pt22k_ft22ktoade20k"""     ):
						UpperCAmelCase_:		int       			=    (1, 1_5_0, 1_6_0, 1_6_0)
						UpperCAmelCase_:		Dict       			=    torch.tensor(
						    [
						        [[-4.3305, -2.3049, -3.0161], [-2.9591, -1.5305, -2.2251], [-3.4198, -1.8004, -2.9062]],
						        [[-5.8922, -3.7435, -4.3978], [-4.2063, -2.7872, -3.4755], [-4.2791, -3.1874, -4.1681]],
						        [[0.9895, 4.3467, 4.7663], [4.2476, 5.6830, 6.1518], [4.5550, 6.2495, 6.5154]],
						    ]     )
			else:
						raise ValueError("""Can't verify logits as model is not supported"""     )

			assert logits.shape == expected_shape, "Shape of logits not as expected"
			if not has_lm_head:
						if is_semantic:
									assert torch.allclose(
									    logits[0, :3, :3, :3]		,	lowerCAmelCase__		,	atol=1e-3     ), "First elements of logits not as expected"
						else:
									print("""Predicted class idx:"""		,	logits.argmax(-1     ).item()     )
									assert torch.allclose(
									    logits[0, :3]		,	lowerCAmelCase__		,	atol=1e-3     ), "First elements of logits not as expected"
									assert logits.argmax(-1     ).item() == expected_class_idx, "Predicted class index not as expected"

			Path(lowerCAmelCase__     ).mkdir(exist_ok=lowerCAmelCase__     )
			print(F'Saving model to {pytorch_dump_folder_path}'     )
			model.save_pretrained(lowerCAmelCase__     )
			print(F'Saving image processor to {pytorch_dump_folder_path}'     )
			image_processor.save_pretrained(lowerCAmelCase__     )


if __name__ == "__main__":
	a			:     Any     = argparse.ArgumentParser()

	parser.add_argument(
	    '--checkpoint_url',
	    default='https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth',
	    type=str,
	    help='URL to the original PyTorch checkpoint (.pth file).',
	)
	parser.add_argument(
	    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the folder to output PyTorch model.'
	)
	a			:     Union[str, Any]     = parser.parse_args()
	convert_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)











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

			def __init__(self,      SCREAMING_SNAKE_CASE_				)					->      Any:
						UpperCAmelCase_:		List[str]       			=    parent
						UpperCAmelCase_:		Tuple       			=    13
						UpperCAmelCase_:		Optional[Any]       			=    7
						UpperCAmelCase_:		List[Any]       			=    True
						UpperCAmelCase_:		List[str]       			=    True
						UpperCAmelCase_:		int       			=    True
						UpperCAmelCase_:		Optional[int]       			=    True
						UpperCAmelCase_:		str       			=    True
						UpperCAmelCase_:		Union[str, Any]       			=    99
						UpperCAmelCase_:		str       			=    32
						UpperCAmelCase_:		str       			=    4
						UpperCAmelCase_:		str       			=    2
						UpperCAmelCase_:		Dict       			=    4
						UpperCAmelCase_:		List[Any]       			=    37
						UpperCAmelCase_:		str       			=    """gelu"""
						UpperCAmelCase_:		int       			=    0.1
						UpperCAmelCase_:		Optional[int]       			=    0.1
						UpperCAmelCase_:		int       			=    512
						UpperCAmelCase_:		List[str]       			=    16
						UpperCAmelCase_:		Optional[int]       			=    2
						UpperCAmelCase_:		List[Any]       			=    0.0_2
						UpperCAmelCase_:		Tuple       			=    3
						UpperCAmelCase_:		Tuple       			=    4
						UpperCAmelCase_:		Union[str, Any]       			=    None
						UpperCAmelCase_:		Dict       			=    self.vocab_size - 1
						UpperCAmelCase_:		Union[str, Any]       			=    self.vocab_size - 1
						UpperCAmelCase_:		Any       			=    self.vocab_size - 1

			def     __snake_case					(self				)					->      Tuple:
						UpperCAmelCase_:		str       			=    ids_tensor([self.batch_size, self.seq_length],      self.vocab_size				)

						UpperCAmelCase_:		Tuple       			=    None
						if self.use_input_mask:
									UpperCAmelCase_:		int       			=    random_attention_mask([self.batch_size, self.seq_length]				)

						UpperCAmelCase_:		Tuple       			=    None
						if self.use_token_type_ids:
									UpperCAmelCase_:		List[Any]       			=    ids_tensor([self.batch_size, self.seq_length],      self.type_vocab_size				)

						UpperCAmelCase_:		Optional[Any]       			=    None
						if self.use_mc_token_ids:
									UpperCAmelCase_:		str       			=    ids_tensor([self.batch_size, self.num_choices],      self.seq_length				)

						UpperCAmelCase_:		int       			=    None
						UpperCAmelCase_:		Optional[int]       			=    None
						UpperCAmelCase_:		Union[str, Any]       			=    None
						if self.use_labels:
									UpperCAmelCase_:		List[str]       			=    ids_tensor([self.batch_size],      self.type_sequence_label_size				)
									UpperCAmelCase_:		str       			=    ids_tensor([self.batch_size, self.seq_length],      self.num_labels				)
									UpperCAmelCase_:		int       			=    ids_tensor([self.batch_size],      self.num_choices				)

						UpperCAmelCase_:		List[str]       			=    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=SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		int       			=    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,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      *SCREAMING_SNAKE_CASE_				)					->      Dict:
						UpperCAmelCase_:		Any       			=    TFGPTJModel(config=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		str       			=    {
						    """input_ids""": input_ids,
						    """attention_mask""": input_mask,
						    """token_type_ids""": token_type_ids,
						}
						UpperCAmelCase_:		str       			=    model(SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Optional[int]       			=    [input_ids, None, input_mask]  # None is the input for 'past'
						UpperCAmelCase_:		List[str]       			=    model(SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Tuple       			=    model(SCREAMING_SNAKE_CASE_				)

						self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.seq_length, self.hidden_size)				)

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      *SCREAMING_SNAKE_CASE_				)					->      Union[str, Any]:
						UpperCAmelCase_:		List[Any]       			=    TFGPTJModel(config=SCREAMING_SNAKE_CASE_				)

						# first forward pass
						UpperCAmelCase_:		Dict       			=    model(SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_,      use_cache=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[Any]       			=    model(SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[Any]       			=    model(SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_,      use_cache=SCREAMING_SNAKE_CASE_				)

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

						UpperCAmelCase_			,       UpperCAmelCase_:		List[Any]       			=    outputs.to_tuple()

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

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

						UpperCAmelCase_:		Tuple       			=    model(SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_				)["""last_hidden_state"""]
						UpperCAmelCase_:		List[str]       			=    model(SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_,      past_key_values=SCREAMING_SNAKE_CASE_				)[
						    """last_hidden_state"""
						]

						# select random slice
						UpperCAmelCase_:		List[Any]       			=    int(ids_tensor((1,),      shape_list(SCREAMING_SNAKE_CASE_				)[-1]				)				)
						UpperCAmelCase_:		Any       			=    output_from_no_past[:, -1, random_slice_idx]
						UpperCAmelCase_:		str       			=    output_from_past[:, 0, random_slice_idx]

						# test that outputs are equal for slice
						tf.debugging.assert_near(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      rtol=1E-6				)

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      *SCREAMING_SNAKE_CASE_				)					->      int:
						UpperCAmelCase_:		List[str]       			=    TFGPTJModel(config=SCREAMING_SNAKE_CASE_				)

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

						# first forward pass
						UpperCAmelCase_			,       UpperCAmelCase_:		Optional[Any]       			=    model(SCREAMING_SNAKE_CASE_,      attention_mask=SCREAMING_SNAKE_CASE_				).to_tuple()

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

						# change a random masked slice from input_ids
						UpperCAmelCase_:		Any       			=    ids_tensor((1,),      SCREAMING_SNAKE_CASE_				).numpy() + 1
						UpperCAmelCase_:		Any       			=    ids_tensor((self.batch_size, self.seq_length),      config.vocab_size				)
						UpperCAmelCase_:		List[Any]       			=    tf.range(self.seq_length				) == (self.seq_length - random_seq_idx_to_change)
						UpperCAmelCase_:		Optional[int]       			=    tf.transpose(
						    tf.broadcast_to(tf.expand_dims(SCREAMING_SNAKE_CASE_,      -1				),      (self.seq_length, self.batch_size)				)				)
						UpperCAmelCase_:		Any       			=    tf.where(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)

						# append to next input_ids and attn_mask
						UpperCAmelCase_:		Any       			=    tf.concat([input_ids, next_tokens],      axis=-1				)
						UpperCAmelCase_:		List[str]       			=    tf.concat([attn_mask, tf.ones((shape_list(SCREAMING_SNAKE_CASE_				)[0], 1),      dtype=tf.intaa				)],      axis=1				)

						# get two different outputs
						UpperCAmelCase_:		List[Any]       			=    model(SCREAMING_SNAKE_CASE_,      attention_mask=SCREAMING_SNAKE_CASE_				)["""last_hidden_state"""]
						UpperCAmelCase_:		Optional[Any]       			=    model(SCREAMING_SNAKE_CASE_,      past_key_values=SCREAMING_SNAKE_CASE_,      attention_mask=SCREAMING_SNAKE_CASE_				)[
						    """last_hidden_state"""
						]

						# select random slice
						UpperCAmelCase_:		Union[str, Any]       			=    int(ids_tensor((1,),      shape_list(SCREAMING_SNAKE_CASE_				)[-1]				)				)
						UpperCAmelCase_:		Tuple       			=    output_from_no_past[:, -1, random_slice_idx]
						UpperCAmelCase_:		int       			=    output_from_past[:, 0, random_slice_idx]

						# test that outputs are equal for slice
						tf.debugging.assert_near(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      rtol=1E-12				)

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      *SCREAMING_SNAKE_CASE_				)					->      Union[str, Any]:
						UpperCAmelCase_:		int       			=    TFGPTJModel(config=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Optional[int]       			=    input_ids[:1, :]
						UpperCAmelCase_:		Tuple       			=    input_mask[:1, :]
						UpperCAmelCase_:		Union[str, Any]       			=    token_type_ids[:1, :]
						UpperCAmelCase_:		Any       			=    1

						# first forward pass
						UpperCAmelCase_:		Optional[Any]       			=    model(SCREAMING_SNAKE_CASE_,      attention_mask=SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_,      use_cache=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_			,       UpperCAmelCase_:		Optional[int]       			=    outputs.to_tuple()

						# create hypothetical next token and extent to next_input_ids
						UpperCAmelCase_:		Dict       			=    ids_tensor((self.batch_size, 3),      config.vocab_size				)
						UpperCAmelCase_:		List[str]       			=    ids_tensor((self.batch_size, 3),      2				)
						UpperCAmelCase_:		Optional[int]       			=    ids_tensor((self.batch_size, 3),      self.type_vocab_size				)

						# append to next input_ids and token_type_ids
						UpperCAmelCase_:		Any       			=    tf.concat([input_ids, next_tokens],      axis=-1				)
						UpperCAmelCase_:		List[Any]       			=    tf.concat([input_mask, next_attn_mask],      axis=-1				)
						UpperCAmelCase_:		Tuple       			=    tf.concat([token_type_ids, next_token_types],      axis=-1				)

						UpperCAmelCase_:		str       			=    model(
						    SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_,      attention_mask=SCREAMING_SNAKE_CASE_				)["""last_hidden_state"""]
						UpperCAmelCase_:		Optional[Any]       			=    model(
						    SCREAMING_SNAKE_CASE_,      token_type_ids=SCREAMING_SNAKE_CASE_,      attention_mask=SCREAMING_SNAKE_CASE_,      past_key_values=SCREAMING_SNAKE_CASE_,      )["""last_hidden_state"""]
						self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]				)

						# select random slice
						UpperCAmelCase_:		List[str]       			=    int(ids_tensor((1,),      shape_list(SCREAMING_SNAKE_CASE_				)[-1]				)				)
						UpperCAmelCase_:		List[str]       			=    output_from_no_past[:, -3:, random_slice_idx]
						UpperCAmelCase_:		List[str]       			=    output_from_past[:, :, random_slice_idx]

						# test that outputs are equal for slice
						tf.debugging.assert_near(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      rtol=1E-3				)

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      *SCREAMING_SNAKE_CASE_				)					->      Optional[int]:
						UpperCAmelCase_:		Tuple       			=    TFGPTJForCausalLM(config=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[int]       			=    {
						    """input_ids""": input_ids,
						    """attention_mask""": input_mask,
						    """token_type_ids""": token_type_ids,
						}
						UpperCAmelCase_:		Union[str, Any]       			=    model(SCREAMING_SNAKE_CASE_				)
						self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.vocab_size)				)

			def     __snake_case					(self				)					->      Optional[Any]:
						UpperCAmelCase_:		Optional[int]       			=    self.prepare_config_and_inputs()

						(
						    (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       (
						    UpperCAmelCase_
						)			,       
						):		Optional[int]       			=    config_and_inputs

						UpperCAmelCase_:		int       			=    {
						    """input_ids""": input_ids,
						    """token_type_ids""": token_type_ids,
						    """attention_mask""": input_mask,
						}
						return config, inputs_dict







@require_tf
class 							_a							(      _lowerCAmelCase   ,		_lowerCAmelCase   ,		_lowerCAmelCase   ,		unittest.TestCase				):
			A					=  (
			    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
			    if is_tf_available()
			    else ()
			)

			A					=  (TFGPTJForCausalLM,) if is_tf_available() else ()
			A					=  (
			    {
			        '''feature-extraction''': TFGPTJModel,
			        '''question-answering''': TFGPTJForQuestionAnswering,
			        '''text-classification''': TFGPTJForSequenceClassification,
			        '''text-generation''': TFGPTJForCausalLM,
			        '''zero-shot''': TFGPTJForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			A					=  False
			A					=  False
			A					=  False
			A					=  False

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)					->      int:
						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				)					->      Tuple:
						UpperCAmelCase_:		Union[str, Any]       			=    TFGPTJModelTester(self				)
						UpperCAmelCase_:		str       			=    ConfigTester(self,      config_class=SCREAMING_SNAKE_CASE_,      n_embd=37				)

			def     __snake_case					(self				)					->      Optional[Any]:
						self.config_tester.run_common_tests()

			def     __snake_case					(self				)					->      Optional[int]:
						UpperCAmelCase_:		Optional[Any]       			=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_gptj_model(*SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self				)					->      Tuple:
						UpperCAmelCase_:		Tuple       			=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_gptj_model_past(*SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self				)					->      Any:
						UpperCAmelCase_:		Dict       			=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_gptj_model_attention_mask_past(*SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self				)					->      Dict:
						UpperCAmelCase_:		List[Any]       			=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_gptj_model_past_large_inputs(*SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self				)					->      str:
						UpperCAmelCase_:		Dict       			=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_gptj_lm_head_model(*SCREAMING_SNAKE_CASE_				)

			@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				)					->      Any:
						UpperCAmelCase_:		List[str]       			=    TFGPTJModel.from_pretrained("""EleutherAI/gpt-j-6B""",      from_pt=SCREAMING_SNAKE_CASE_				)
						self.assertIsNotNone(SCREAMING_SNAKE_CASE_				)

			@unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""				)
			def     __snake_case					(self				)					->      Tuple:
						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 							_a							(      unittest.TestCase				):

			def     __snake_case					(self				)					->      str:
						UpperCAmelCase_:		List[Any]       			=    TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B""",      from_pt=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		List[str]       			=    tf.convert_to_tensor([[464, 3290]],      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_:		str       			=    [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
						# fmt: on
						UpperCAmelCase_:		Tuple       			=    model.generate(SCREAMING_SNAKE_CASE_,      do_sample=SCREAMING_SNAKE_CASE_				)
						self.assertListEqual(output_ids[0].numpy().tolist(),      SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self				)					->      Any:
						UpperCAmelCase_:		Optional[int]       			=    AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B""",      revision="""float16"""				)
						UpperCAmelCase_:		Dict       			=    TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B""",      revision="""float16""",      from_pt=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Tuple       			=    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_:		List[Any]       			=    model.generate(**SCREAMING_SNAKE_CASE_,      do_sample=SCREAMING_SNAKE_CASE_,      seed=[42, 0]				)
						UpperCAmelCase_:		Dict       			=    tokenizer.decode(output_ids[0],      skip_special_tokens=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Union[str, Any]       			=    """Today is a nice day and I’m going to go for a walk. I’"""
						self.assertEqual(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self				)					->      List[Any]:
						UpperCAmelCase_:		List[Any]       			=    TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B""",      revision="""float16""",      from_pt=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		int       			=    AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B""",      revision="""float16"""				)

						UpperCAmelCase_:		int       			=    """left"""

						# Define PAD Token = EOS Token = 50256
						UpperCAmelCase_:		int       			=    tokenizer.eos_token
						UpperCAmelCase_:		Optional[int]       			=    model.config.eos_token_id

						# use different length sentences to test batching
						UpperCAmelCase_:		Any       			=    [
						    """Hello, my dog is a little""",
						    """Today, I""",
						]
						UpperCAmelCase_:		Any       			=    [
						    """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				)					->      Any:
						# Confirms that we get the expected results with left-padded beam search
						UpperCAmelCase_			,       UpperCAmelCase_			,       UpperCAmelCase_			,       UpperCAmelCase_:		Optional[Any]       			=    self._get_beam_search_test_objects()

						UpperCAmelCase_:		Optional[int]       			=    tokenizer(SCREAMING_SNAKE_CASE_,      return_tensors="""tf""",      padding=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Dict       			=    model.generate(**SCREAMING_SNAKE_CASE_,      do_sample=SCREAMING_SNAKE_CASE_,      num_beams=2				)
						UpperCAmelCase_:		Optional[Any]       			=    tokenizer.batch_decode(SCREAMING_SNAKE_CASE_,      skip_special_tokens=SCREAMING_SNAKE_CASE_				)
						self.assertListEqual(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self				)					->      List[Any]:
						# Confirms that left-padding is working properly
						UpperCAmelCase_			,       UpperCAmelCase_			,       UpperCAmelCase_			,       UpperCAmelCase_:		List[Any]       			=    self._get_beam_search_test_objects()

						UpperCAmelCase_:		List[Any]       			=    tokenizer(SCREAMING_SNAKE_CASE_,      return_tensors="""tf""",      padding=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Dict       			=    tokenizer(sentences[0],      return_tensors="""tf"""				)
						UpperCAmelCase_:		Optional[Any]       			=    model.generate(**SCREAMING_SNAKE_CASE_,      do_sample=SCREAMING_SNAKE_CASE_,      num_beams=2				)
						UpperCAmelCase_:		Optional[int]       			=    (
						    shape_list(inputs_non_padded["""input_ids"""]				)[-1]
						    - tf.reduce_sum(tf.cast(inputs["""attention_mask"""][-1],      tf.intaa				)				).numpy()
						)
						UpperCAmelCase_:		List[str]       			=    tokenizer(sentences[1],      return_tensors="""tf"""				)
						UpperCAmelCase_:		Any       			=    model.generate(
						    **SCREAMING_SNAKE_CASE_,      do_sample=SCREAMING_SNAKE_CASE_,      num_beams=2,      max_length=model.config.max_length - num_paddings				)
						UpperCAmelCase_:		Dict       			=    tokenizer.decode(output_non_padded[0],      skip_special_tokens=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		str       			=    tokenizer.decode(output_padded[0],      skip_special_tokens=SCREAMING_SNAKE_CASE_				)
						self.assertListEqual(SCREAMING_SNAKE_CASE_,      [non_padded_sentence, padded_sentence]				)

			def     __snake_case					(self				)					->      str:
						# Confirms that XLA is working properly
						UpperCAmelCase_			,       UpperCAmelCase_			,       UpperCAmelCase_			,       UpperCAmelCase_:		Optional[int]       			=    self._get_beam_search_test_objects()

						UpperCAmelCase_:		Tuple       			=    tokenizer(SCREAMING_SNAKE_CASE_,      return_tensors="""tf""",      padding=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		int       			=    tf.function(model.generate,      jit_compile=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Union[str, Any]       			=    xla_generate(**SCREAMING_SNAKE_CASE_,      do_sample=SCREAMING_SNAKE_CASE_,      num_beams=2				)
						UpperCAmelCase_:		Any       			=    tokenizer.batch_decode(SCREAMING_SNAKE_CASE_,      skip_special_tokens=SCREAMING_SNAKE_CASE_				)
						self.assertListEqual(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)






