


from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import SwinConfig
from transformers.testing_utils import require_tf, require_vision, slow, to_atuple
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers.models.swin.modeling_tf_swin import (
				    TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFSwinForImageClassification,
				    TFSwinForMaskedImageModeling,
				    TFSwinModel,
				)


if is_vision_available():
				from PIL import Image

				from transformers import AutoImageProcessor







class      __lowercase  :





							"""simple docstring"""







							def __init__(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__=1_3       ,		lowerCAmelCase__=3_2       ,		lowerCAmelCase__=2       ,		lowerCAmelCase__=3       ,		lowerCAmelCase__=1_6       ,		lowerCAmelCase__=[1, 2, 1]       ,		lowerCAmelCase__=[2, 2, 4]       ,		lowerCAmelCase__=2       ,		lowerCAmelCase__=2.0       ,		lowerCAmelCase__=True       ,		lowerCAmelCase__=0.0       ,		lowerCAmelCase__=0.0       ,		lowerCAmelCase__=0.1       ,		lowerCAmelCase__="gelu"       ,		lowerCAmelCase__=False       ,		lowerCAmelCase__=True       ,		lowerCAmelCase__=0.02       ,		lowerCAmelCase__=1E-5       ,		lowerCAmelCase__=True       ,		lowerCAmelCase__=None       ,		lowerCAmelCase__=True       ,		lowerCAmelCase__=1_0       ,		lowerCAmelCase__=8       ,		):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			str       						= parent
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= batch_size
										SCREAMING_SNAKE_CASE_						:			Tuple       						= image_size
										SCREAMING_SNAKE_CASE_						:			Tuple       						= patch_size
										SCREAMING_SNAKE_CASE_						:			List[str]       						= num_channels
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= embed_dim
										SCREAMING_SNAKE_CASE_						:			Any       						= depths
										SCREAMING_SNAKE_CASE_						:			Dict       						= num_heads
										SCREAMING_SNAKE_CASE_						:			int       						= window_size
										SCREAMING_SNAKE_CASE_						:			Dict       						= mlp_ratio
										SCREAMING_SNAKE_CASE_						:			str       						= qkv_bias
										SCREAMING_SNAKE_CASE_						:			List[str]       						= hidden_dropout_prob
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= attention_probs_dropout_prob
										SCREAMING_SNAKE_CASE_						:			str       						= drop_path_rate
										SCREAMING_SNAKE_CASE_						:			str       						= hidden_act
										SCREAMING_SNAKE_CASE_						:			Dict       						= use_absolute_embeddings
										SCREAMING_SNAKE_CASE_						:			Dict       						= patch_norm
										SCREAMING_SNAKE_CASE_						:			Any       						= layer_norm_eps
										SCREAMING_SNAKE_CASE_						:			Any       						= initializer_range
										SCREAMING_SNAKE_CASE_						:			Dict       						= is_training
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= scope
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= use_labels
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= type_sequence_label_size
										SCREAMING_SNAKE_CASE_						:			int       						= encoder_stride






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Any       						= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= None
										if self.use_labels:
													SCREAMING_SNAKE_CASE_						:			List[str]       						= ids_tensor([self.batch_size]       ,		self.type_sequence_label_size    )

										SCREAMING_SNAKE_CASE_						:			int       						= self.get_config()

										return config, pixel_values, labels






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										return SwinConfig(
										    image_size=self.image_size       ,		patch_size=self.patch_size       ,		num_channels=self.num_channels       ,		embed_dim=self.embed_dim       ,		depths=self.depths       ,		num_heads=self.num_heads       ,		window_size=self.window_size       ,		mlp_ratio=self.mlp_ratio       ,		qkv_bias=self.qkv_bias       ,		hidden_dropout_prob=self.hidden_dropout_prob       ,		attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,		drop_path_rate=self.drop_path_rate       ,		hidden_act=self.hidden_act       ,		use_absolute_embeddings=self.use_absolute_embeddings       ,		path_norm=self.patch_norm       ,		layer_norm_eps=self.layer_norm_eps       ,		initializer_range=self.initializer_range       ,		encoder_stride=self.encoder_stride       ,		)






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			str       						= TFSwinModel(config=lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			str       						= model(lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths    ) - 1))
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= int(config.embed_dim * 2 ** (len(config.depths    ) - 1)    )

										self.parent.assertEqual(result.last_hidden_state.shape       ,		(self.batch_size, expected_seq_len, expected_dim)    )






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= TFSwinForMaskedImageModeling(config=lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Any       						= model(lowerCAmelCase__    )
										self.parent.assertEqual(
										    result.logits.shape       ,		(self.batch_size, self.num_channels, self.image_size, self.image_size)    )

										# test greyscale images
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= 1
										SCREAMING_SNAKE_CASE_						:			str       						= TFSwinForMaskedImageModeling(lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			List[str]       						= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
										SCREAMING_SNAKE_CASE_						:			int       						= model(lowerCAmelCase__    )
										self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, 1, self.image_size, self.image_size)    )






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			str       						= self.type_sequence_label_size
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= TFSwinForImageClassification(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= model(lowerCAmelCase__       ,		labels=lowerCAmelCase__    )
										self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.type_sequence_label_size)    )

										# test greyscale images
										SCREAMING_SNAKE_CASE_						:			int       						= 1
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= TFSwinForImageClassification(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= model(lowerCAmelCase__    )
										self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.type_sequence_label_size)    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= self.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			str       						= config_and_inputs
										SCREAMING_SNAKE_CASE_						:			Any       						= {'pixel_values': pixel_values}
										return config, inputs_dict







@require_tf
class      __lowercase  (__SCREAMING_SNAKE_CASE     ,    __SCREAMING_SNAKE_CASE     ,    unittest.TestCase   ):





							"""simple docstring"""







							_UpperCAmelCase     =		(
							    (
							        TFSwinModel,
							        TFSwinForImageClassification,
							        TFSwinForMaskedImageModeling,
							    )
							    if is_tf_available()
							    else ()
							)
							_UpperCAmelCase     =		(
							    {"""feature-extraction""": TFSwinModel, """image-classification""": TFSwinForImageClassification}
							    if is_tf_available()
							    else {}
							)

							_UpperCAmelCase     =		False
							_UpperCAmelCase     =		False
							_UpperCAmelCase     =		False
							_UpperCAmelCase     =		False






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= TFSwinModelTester(self    )
										SCREAMING_SNAKE_CASE_						:			Dict       						= ConfigTester(self       ,		config_class=lowerCAmelCase__       ,		embed_dim=3_7    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										self.create_and_test_config_common_properties()
										self.config_tester.create_and_test_config_to_json_string()
										self.config_tester.create_and_test_config_to_json_file()
										self.config_tester.create_and_test_config_from_and_save_pretrained()
										self.config_tester.create_and_test_config_with_num_labels()
										self.config_tester.check_config_can_be_init_without_params()
										self.config_tester.check_config_arguments_init()






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										return






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			List[str]       						= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			List[str]       						= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_image_modeling(*lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			List[str]       						= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase__    )






							@unittest.skip(reason='Swin does not use inputs_embeds'    )
							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										pass






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			List[Any]       						= self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													SCREAMING_SNAKE_CASE_						:			List[Any]       						= model_class(lowerCAmelCase__    )
													self.assertIsInstance(model.get_input_embeddings()       ,		tf.keras.layers.Layer    )
													SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(lowerCAmelCase__       ,		tf.keras.layers.Dense    )    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			List[str]       						= self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													SCREAMING_SNAKE_CASE_						:			str       						= model_class(lowerCAmelCase__    )
													SCREAMING_SNAKE_CASE_						:			str       						= inspect.signature(model.call    )
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													SCREAMING_SNAKE_CASE_						:			int       						= [*signature.parameters.keys()]

													SCREAMING_SNAKE_CASE_						:			Dict       						= ['pixel_values']
													self.assertListEqual(arg_names[:1]       ,		lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			List[str]       						= self.model_tester.prepare_config_and_inputs_for_common()
										SCREAMING_SNAKE_CASE_						:			Dict       						= True

										for model_class in self.all_model_classes:
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= True
													SCREAMING_SNAKE_CASE_						:			str       						= False
													SCREAMING_SNAKE_CASE_						:			Dict       						= True
													SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= model_class(lowerCAmelCase__    )
													SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= model(**self._prepare_for_class(lowerCAmelCase__       ,		lowerCAmelCase__    )    )
													SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= outputs.attentions
													SCREAMING_SNAKE_CASE_						:			int       						= len(self.model_tester.depths    )
													self.assertEqual(len(lowerCAmelCase__    )       ,		lowerCAmelCase__    )

													# check that output_attentions also work using config
													del inputs_dict["output_attentions"]
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= True
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= config.window_size**2
													SCREAMING_SNAKE_CASE_						:			List[str]       						= model_class(lowerCAmelCase__    )
													SCREAMING_SNAKE_CASE_						:			str       						= model(**self._prepare_for_class(lowerCAmelCase__       ,		lowerCAmelCase__    )    )
													SCREAMING_SNAKE_CASE_						:			Dict       						= outputs.attentions
													self.assertEqual(len(lowerCAmelCase__    )       ,		lowerCAmelCase__    )

													self.assertListEqual(
													    list(attentions[0].shape[-3:]    )       ,		[self.model_tester.num_heads[0], window_size_squared, window_size_squared]       ,		)
													SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= len(lowerCAmelCase__    )

													# Check attention is always last and order is fine
													SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= True
													SCREAMING_SNAKE_CASE_						:			Dict       						= True
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= model_class(lowerCAmelCase__    )
													SCREAMING_SNAKE_CASE_						:			List[Any]       						= model(**self._prepare_for_class(lowerCAmelCase__       ,		lowerCAmelCase__    )    )

													if hasattr(self.model_tester       ,		'num_hidden_states_types'    ):
																SCREAMING_SNAKE_CASE_						:			int       						= self.model_tester.num_hidden_states_types
													else:
																# also another +1 for reshaped_hidden_states
																SCREAMING_SNAKE_CASE_						:			int       						= 2
													self.assertEqual(out_len + added_hidden_states       ,		len(lowerCAmelCase__    )    )

													SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= outputs.attentions

													self.assertEqual(len(lowerCAmelCase__    )       ,		lowerCAmelCase__    )

													self.assertListEqual(
													    list(self_attentions[0].shape[-3:]    )       ,		[self.model_tester.num_heads[0], window_size_squared, window_size_squared]       ,		)






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Tuple       						= model_class(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= model(**self._prepare_for_class(lowerCAmelCase__       ,		lowerCAmelCase__    )    )
										SCREAMING_SNAKE_CASE_						:			str       						= outputs.hidden_states

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= getattr(
										    self.model_tester       ,		'expected_num_hidden_layers'       ,		len(self.model_tester.depths    ) + 1    )
										self.assertEqual(len(lowerCAmelCase__    )       ,		lowerCAmelCase__    )

										# Swin has a different seq_length
										SCREAMING_SNAKE_CASE_						:			str       						= to_atuple(config.patch_size    )

										SCREAMING_SNAKE_CASE_						:			Tuple       						= (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

										self.assertListEqual(
										    list(hidden_states[0].shape[-2:]    )       ,		[num_patches, self.model_tester.embed_dim]       ,		)

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= outputs.reshaped_hidden_states
										self.assertEqual(len(lowerCAmelCase__    )       ,		lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			str       						= reshaped_hidden_states[0].shape

										SCREAMING_SNAKE_CASE_						:			int       						= tf.reshape(reshaped_hidden_states[0]       ,		(batch_size, num_channels, height * width)    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= tf.transpose(lowerCAmelCase__       ,		(0, 2, 1)    )

										self.assertListEqual(
										    list(reshaped_hidden_states.shape[-2:]    )       ,		[num_patches, self.model_tester.embed_dim]       ,		)






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			Tuple       						= self.model_tester.prepare_config_and_inputs_for_common()

										SCREAMING_SNAKE_CASE_						:			List[str]       						= to_atuple(self.model_tester.image_size    )

										for model_class in self.all_model_classes:
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= True
													self.check_hidden_states_output(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    )

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													SCREAMING_SNAKE_CASE_						:			Dict       						= True

													self.check_hidden_states_output(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			List[Any]       						= self.model_tester.prepare_config_and_inputs_for_common()
										SCREAMING_SNAKE_CASE_						:			int       						= 3

										SCREAMING_SNAKE_CASE_						:			Any       						= to_atuple(self.model_tester.image_size    )
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= to_atuple(config.patch_size    )

										SCREAMING_SNAKE_CASE_						:			str       						= image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

										for model_class in self.all_model_classes:
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= True
													self.check_hidden_states_output(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		(padded_height, padded_width)    )

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													SCREAMING_SNAKE_CASE_						:			List[str]       						= True
													self.check_hidden_states_output(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		(padded_height, padded_width)    )






							@slow
							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										for model_name in TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE_						:			Any       						= TFSwinModel.from_pretrained(lowerCAmelCase__    )
													self.assertIsNotNone(lowerCAmelCase__    )







@require_vision
@require_tf
class      __lowercase  (unittest.TestCase   ):





							"""simple docstring"""







							@cached_property
							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										return (
										    AutoImageProcessor.from_pretrained('microsoft/swin-tiny-patch4-window7-224'    )
										    if is_vision_available()
										    else None
										)






							@slow
							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			int       						= TFSwinForImageClassification.from_pretrained('microsoft/swin-tiny-patch4-window7-224'    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= self.default_image_processor

										SCREAMING_SNAKE_CASE_						:			int       						= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'    )
										SCREAMING_SNAKE_CASE_						:			str       						= image_processor(images=lowerCAmelCase__       ,		return_tensors='tf'    )

										# forward pass
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= model(lowerCAmelCase__    )

										# verify the logits
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= tf.TensorShape((1, 1_0_0_0)    )
										self.assertEqual(outputs.logits.shape       ,		lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= tf.constant([-0.0_948, -0.6_454, -0.0_921]    )
										self.assertTrue(np.allclose(outputs.logits[0, :3]       ,		lowerCAmelCase__       ,		atol=1E-4    )    )









import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
				import os

				# The slow tests are often failing with OOM error on GPU
				# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
				# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
				lowerCAmelCase__       :	Dict     				='platform'

				import jax
				import jax.numpy as jnp

				from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer






def 		a__       (							A__,	A__,	A__=None,	A__=None							):
			if attention_mask is None:
						SCREAMING_SNAKE_CASE_						:			Tuple       						= np.where(input_ids != config.pad_token_id,	1,	0							)
			return {
			    "input_ids": input_ids,
			    "attention_mask": attention_mask,
			}







@require_flax
class      __lowercase  :





							"""simple docstring"""







							def __init__(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__=1_3       ,		lowerCAmelCase__=7       ,		lowerCAmelCase__=True       ,		lowerCAmelCase__=False       ,		lowerCAmelCase__=9_9       ,		lowerCAmelCase__=1_6       ,		lowerCAmelCase__=2       ,		lowerCAmelCase__=4       ,		lowerCAmelCase__=4       ,		lowerCAmelCase__="gelu"       ,		lowerCAmelCase__=0.1       ,		lowerCAmelCase__=0.1       ,		lowerCAmelCase__=2_0       ,		lowerCAmelCase__=2       ,		lowerCAmelCase__=1       ,		lowerCAmelCase__=0       ,		lowerCAmelCase__=1_6       ,		lowerCAmelCase__=1_6       ,		lowerCAmelCase__=0.02       ,		):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			str       						= parent
										SCREAMING_SNAKE_CASE_						:			List[str]       						= batch_size
										SCREAMING_SNAKE_CASE_						:			str       						= seq_length
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= is_training
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= use_labels
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= vocab_size
										SCREAMING_SNAKE_CASE_						:			Tuple       						= hidden_size
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= num_hidden_layers
										SCREAMING_SNAKE_CASE_						:			Dict       						= num_attention_heads
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= intermediate_size
										SCREAMING_SNAKE_CASE_						:			Tuple       						= hidden_act
										SCREAMING_SNAKE_CASE_						:			Dict       						= hidden_dropout_prob
										SCREAMING_SNAKE_CASE_						:			int       						= attention_probs_dropout_prob
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= max_position_embeddings
										SCREAMING_SNAKE_CASE_						:			Tuple       						= eos_token_id
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= pad_token_id
										SCREAMING_SNAKE_CASE_						:			str       						= bos_token_id
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= embed_dim
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= word_embed_proj_dim
										SCREAMING_SNAKE_CASE_						:			Tuple       						= initializer_range
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= False






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Dict       						= np.clip(ids_tensor([self.batch_size, self.seq_length - 1]       ,		self.vocab_size    )       ,		3       ,		self.vocab_size    )
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)       ,		dtype=np.intaa    ))       ,		-1    )

										SCREAMING_SNAKE_CASE_						:			List[str]       						= OPTConfig(
										    vocab_size=self.vocab_size       ,		hidden_size=self.hidden_size       ,		num_hidden_layers=self.num_hidden_layers       ,		num_attention_heads=self.num_attention_heads       ,		ffn_dim=self.intermediate_size       ,		dropout=self.hidden_dropout_prob       ,		attention_dropout=self.attention_probs_dropout_prob       ,		max_position_embeddings=self.max_position_embeddings       ,		eos_token_id=self.eos_token_id       ,		bos_token_id=self.bos_token_id       ,		pad_token_id=self.pad_token_id       ,		embed_dim=self.embed_dim       ,		is_encoder_decoder=lowerCAmelCase__       ,		word_embed_proj_dim=self.word_embed_proj_dim       ,		initializer_range=self.initializer_range       ,		use_cache=lowerCAmelCase__       ,		)
										SCREAMING_SNAKE_CASE_						:			Dict       						= prepare_opt_inputs_dict(lowerCAmelCase__       ,		lowerCAmelCase__    )
										return config, inputs_dict






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			Any       						= self.prepare_config_and_inputs()
										return config, inputs_dict






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			List[Any]       						= 2_0
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= model_class_name(lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			Tuple       						= inputs_dict['input_ids']
										SCREAMING_SNAKE_CASE_						:			Dict       						= inputs_dict['attention_mask']

										SCREAMING_SNAKE_CASE_						:			List[Any]       						= model.init_cache(input_ids.shape[0]       ,		lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= jnp.ones((input_ids.shape[0], max_length)       ,		dtype='i4'    )

										SCREAMING_SNAKE_CASE_						:			List[Any]       						= jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1    )[None, :]       ,		(input_ids.shape[0], input_ids.shape[-1] - 1)       ,		)
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= model(
										    input_ids[:, :-1]       ,		attention_mask=lowerCAmelCase__       ,		past_key_values=lowerCAmelCase__       ,		position_ids=lowerCAmelCase__       ,		)

										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,		dtype='i4'    )
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= model(
										    input_ids[:, -1:]       ,		attention_mask=lowerCAmelCase__       ,		past_key_values=outputs_cache.past_key_values       ,		position_ids=lowerCAmelCase__       ,		)

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= model(lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			List[Any]       						= np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])    )    )
										self.parent.assertTrue(diff < 1E-3       ,		msg=F'''Max diff is {diff}'''    )






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= 2_0
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= model_class_name(lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			Any       						= (
										    inputs_dict['input_ids'],
										    inputs_dict['attention_mask'],
										)

										SCREAMING_SNAKE_CASE_						:			Any       						= jnp.concatenate(
										    [
										        attention_mask,
										        jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])    ),
										    ]       ,		axis=-1       ,		)

										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= model.init_cache(input_ids.shape[0]       ,		lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1    )[None, :]       ,		(input_ids.shape[0], input_ids.shape[-1] - 1)       ,		)

										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= model(
										    input_ids[:, :-1]       ,		attention_mask=lowerCAmelCase__       ,		past_key_values=lowerCAmelCase__       ,		position_ids=lowerCAmelCase__       ,		)
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,		dtype='i4'    )
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= model(
										    input_ids[:, -1:]       ,		past_key_values=outputs_cache.past_key_values       ,		attention_mask=lowerCAmelCase__       ,		position_ids=lowerCAmelCase__       ,		)

										SCREAMING_SNAKE_CASE_						:			Dict       						= model(lowerCAmelCase__       ,		attention_mask=lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			Any       						= np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])    )    )
										self.parent.assertTrue(diff < 1E-3       ,		msg=F'''Max diff is {diff}'''    )







@require_flax
class      __lowercase  (__SCREAMING_SNAKE_CASE     ,    unittest.TestCase     ,    __SCREAMING_SNAKE_CASE   ):





							"""simple docstring"""







							_UpperCAmelCase     =		(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
							_UpperCAmelCase     =		() if is_flax_available() else ()






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= FlaxOPTModelTester(self    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			Any       						= self.model_tester.prepare_config_and_inputs()
										for model_class in self.all_model_classes:
													self.model_tester.check_use_cache_forward(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_						:			Dict       						= self.model_tester.prepare_config_and_inputs()
										for model_class in self.all_model_classes:
													self.model_tester.check_use_cache_forward_with_attn_mask(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    )






							@slow
							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										for model_class_name in self.all_model_classes:
													SCREAMING_SNAKE_CASE_						:			str       						= model_class_name.from_pretrained('facebook/opt-125m'    )
													SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= np.ones((1, 1)    ) * model.config.eos_token_id
													SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= model(lowerCAmelCase__    )
													self.assertIsNotNone(lowerCAmelCase__    )







@require_sentencepiece
@require_flax
class      __lowercase  (unittest.TestCase   ):





							"""simple docstring"""







							@slow
							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			str       						= FlaxOPTModel.from_pretrained('facebook/opt-350m'    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= jnp.array([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]    )
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= model(input_ids=lowerCAmelCase__    ).last_hidden_state
										SCREAMING_SNAKE_CASE_						:			int       						= (1, 1_1, 5_1_2)
										self.assertEqual(output.shape       ,		lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= jnp.array(
										    [[-0.2_867, -1.9_256, -0.3_062], [-1.2_711, -0.1_337, -0.1_897], [0.4_109, 0.1_187, -1.3_142]]    )
										self.assertTrue(jnp.allclose(output[:, :3, :3]       ,		lowerCAmelCase__       ,		atol=4E-2    )    )







@require_flax
@slow
class      __lowercase  (unittest.TestCase   ):





							"""simple docstring"""







							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										super().setUp()
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= 'facebook/opt-350m'






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= FlaxOPTForCausalLM.from_pretrained(self.path_model    )
										SCREAMING_SNAKE_CASE_						:			int       						= GPTaTokenizer.from_pretrained(self.path_model    )

										SCREAMING_SNAKE_CASE_						:			int       						= [
										    'Today is a beautiful day and I want to',
										    'In the city of',
										    'Paris is the capital of France and',
										    'Computers and mobile phones have taken',
										]
										# verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
										SCREAMING_SNAKE_CASE_						:			Any       						= tokenizer(lowerCAmelCase__       ,		return_tensors='jax'       ,		padding=lowerCAmelCase__       ,		add_special_tokens=lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Any       						= model(inputs.input_ids       ,		attention_mask=inputs.attention_mask    )[0].mean(axis=-1    )
										SCREAMING_SNAKE_CASE_						:			Any       						= jnp.array(
										    [
										        [1.3_851, -13.8_923, -10.5_229, -10.7_533, -0.2_309, -10.2_384, -0.5_365, -9.0_947, -5.1_670],
										        [-4.7_073, -10.6_276, -3.9_415, -21.5_242, -0.2_822, -0.2_822, -0.2_822, -0.2_822, -0.2_822],
										        [0.6_247, -3.4_229, -8.9_179, -1.4_297, -14.1_650, 1.4_146, -9.0_218, -0.2_703, -0.2_703],
										        [6.4_783, -1.9_913, -10.7_926, -2.3_336, 1.5_092, -0.9_974, -6.8_213, 1.3_477, 1.3_477],
										    ]    )
										self.assertTrue(jnp.allclose(lowerCAmelCase__       ,		lowerCAmelCase__       ,		atol=4E-2    )    )

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= jax.jit(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= model(inputs.input_ids       ,		attention_mask=inputs.attention_mask    )[0].mean(axis=-1    )
										self.assertTrue(jnp.allclose(lowerCAmelCase__       ,		lowerCAmelCase__       ,		atol=4E-2    )    )







@require_flax
@slow
class      __lowercase  (unittest.TestCase   ):





							"""simple docstring"""







							@property
							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										return [
										    "Today is a beautiful day and I want",
										    "In the city of",
										    "Paris is the capital of France and",
										    "Computers and mobile phones have taken",
										]






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= 'facebook/opt-125m'

										SCREAMING_SNAKE_CASE_						:			List[Any]       						= [
										    'Today is a beautiful day and I want to',
										    'In the city of New York, the city',
										    'Paris is the capital of France and the capital',
										    'Computers and mobile phones have taken over the',
										]

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= []

										SCREAMING_SNAKE_CASE_						:			int       						= FlaxOPTForCausalLM.from_pretrained(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Dict       						= GPTaTokenizer.from_pretrained(lowerCAmelCase__    )

										for prompt in self.prompts:
													SCREAMING_SNAKE_CASE_						:			Any       						= tokenizer(lowerCAmelCase__       ,		return_tensors='jax'    ).input_ids

													SCREAMING_SNAKE_CASE_						:			List[Any]       						= model.generate(lowerCAmelCase__       ,		max_length=1_0    )
													SCREAMING_SNAKE_CASE_						:			int       						= generated_ids[0]

													SCREAMING_SNAKE_CASE_						:			int       						= tokenizer.batch_decode(lowerCAmelCase__       ,		skip_special_tokens=lowerCAmelCase__    )
													predicted_outputs += generated_string

										self.assertListEqual(lowerCAmelCase__       ,		lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Any       						= 'facebook/opt-350m'

										SCREAMING_SNAKE_CASE_						:			List[str]       						= [
										    'Today is a beautiful day and I want to',
										    'In the city of San Francisco, the city',
										    'Paris is the capital of France and the capital',
										    'Computers and mobile phones have taken over the',
										]

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= []
										SCREAMING_SNAKE_CASE_						:			List[str]       						= FlaxOPTForCausalLM.from_pretrained(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Tuple       						= GPTaTokenizer.from_pretrained(lowerCAmelCase__    )

										for prompt in self.prompts:
													SCREAMING_SNAKE_CASE_						:			List[Any]       						= tokenizer(lowerCAmelCase__       ,		return_tensors='jax'    ).input_ids

													SCREAMING_SNAKE_CASE_						:			List[Any]       						= model.generate(lowerCAmelCase__       ,		max_length=1_0    )
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= generated_ids[0]

													SCREAMING_SNAKE_CASE_						:			List[Any]       						= tokenizer.batch_decode(lowerCAmelCase__       ,		skip_special_tokens=lowerCAmelCase__    )
													predicted_outputs += generated_string

										self.assertListEqual(lowerCAmelCase__       ,		lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Any       						= 'facebook/opt-125m'
										SCREAMING_SNAKE_CASE_						:			str       						= [
										    'Today is a beautiful day and I want to thank',
										    'In the city of Rome Canaver Canaver Canaver Canaver',
										]
										SCREAMING_SNAKE_CASE_						:			Any       						= FlaxOPTForCausalLM.from_pretrained(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= GPTaTokenizer.from_pretrained(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= tokenizer(
										    [
										        'Today is a beautiful day and I want to',
										        'In the city of',
										    ]       ,		return_tensors='jax'       ,		padding=lowerCAmelCase__       ,		)

										SCREAMING_SNAKE_CASE_						:			List[str]       						= jax.jit(model.generate    )

										SCREAMING_SNAKE_CASE_						:			List[Any]       						= jit_generate(inputs['input_ids']       ,		attention_mask=inputs['attention_mask']    ).sequences

										SCREAMING_SNAKE_CASE_						:			str       						= tokenizer.batch_decode(lowerCAmelCase__       ,		skip_special_tokens=lowerCAmelCase__    )

										self.assertIsNotNone(lowerCAmelCase__       ,		lowerCAmelCase__    )






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Tuple       						= 'facebook/opt-350m'

										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= GPTaTokenizer.from_pretrained(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			str       						= FlaxOPTForCausalLM.from_pretrained(lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= 'left'

										# use different length sentences to test batching
										SCREAMING_SNAKE_CASE_						:			Dict       						= [
										    'Hello, my dog is a little',
										    'Today, I',
										]

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= tokenizer(lowerCAmelCase__       ,		return_tensors='jax'       ,		padding=lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			str       						= inputs['input_ids']

										SCREAMING_SNAKE_CASE_						:			str       						= model.generate(input_ids=lowerCAmelCase__       ,		attention_mask=inputs['attention_mask']       ,		trace=lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			Dict       						= tokenizer(sentences[0]       ,		return_tensors='jax'    ).input_ids
										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= model.generate(input_ids=lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= inputs_non_padded.shape[-1] - inputs['attention_mask'][-1].sum()
										SCREAMING_SNAKE_CASE_						:			str       						= tokenizer(sentences[1]       ,		return_tensors='jax'    ).input_ids
										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= model.generate(input_ids=lowerCAmelCase__       ,		max_length=model.config.max_length - num_paddings    )

										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= tokenizer.batch_decode(outputs[0]       ,		skip_special_tokens=lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			int       						= tokenizer.decode(output_non_padded[0][0]       ,		skip_special_tokens=lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			str       						= tokenizer.decode(output_padded[0][0]       ,		skip_special_tokens=lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			Optional[int]       						= [
										    'Hello, my dog is a little bit of a dork.\nI\'m a little bit',
										    'Today, I was in the middle of a conversation with a friend about the',
										]
										self.assertListEqual(lowerCAmelCase__       ,		lowerCAmelCase__    )
										self.assertListEqual(lowerCAmelCase__       ,		[non_padded_sentence, padded_sentence]    )





