






import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

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


if is_flax_available():
				import jax
				import jax.numpy as jnp

				from transformers.modeling_flax_pytorch_utils import (
				    convert_pytorch_state_dict_to_flax,
				    load_flax_weights_in_pytorch_model,
				)
				from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
				import torch


class 				__a				:





		'''simple docstring'''



		def __init__(    self			,     UpperCamelCase__			,     UpperCamelCase__=14			,     UpperCamelCase__=7			,     UpperCamelCase__=True			,     UpperCamelCase__=True			,     UpperCamelCase__=False			,     UpperCamelCase__=True			,     UpperCamelCase__=99			,     UpperCamelCase__=32			,     UpperCamelCase__=5			,     UpperCamelCase__=4			,     UpperCamelCase__=37			,     UpperCamelCase__="gelu"			,     UpperCamelCase__=0.1			,     UpperCamelCase__=0.1			,     UpperCamelCase__=512			,     UpperCamelCase__=0.02			,     ):
						SCREAMING_SNAKE_CASE_ :						List[str]    				= parent
						SCREAMING_SNAKE_CASE_ :						int    				= batch_size
						SCREAMING_SNAKE_CASE_ :						List[Any]    				= seq_length
						SCREAMING_SNAKE_CASE_ :						Dict    				= is_training
						SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= use_input_mask
						SCREAMING_SNAKE_CASE_ :						Dict    				= use_token_type_ids
						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= use_labels
						SCREAMING_SNAKE_CASE_ :						str    				= vocab_size
						SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= hidden_size
						SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= num_hidden_layers
						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= num_attention_heads
						SCREAMING_SNAKE_CASE_ :						Dict    				= intermediate_size
						SCREAMING_SNAKE_CASE_ :						Tuple    				= hidden_act
						SCREAMING_SNAKE_CASE_ :						str    				= hidden_dropout_prob
						SCREAMING_SNAKE_CASE_ :						Any    				= attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= max_position_embeddings
						SCREAMING_SNAKE_CASE_ :						str    				= initializer_range
						SCREAMING_SNAKE_CASE_ :						Any    				= None
						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= vocab_size - 1
						SCREAMING_SNAKE_CASE_ :						Tuple    				= vocab_size - 1
						SCREAMING_SNAKE_CASE_ :						List[str]    				= vocab_size - 1


		def 	__snake_case      (    self							):
						SCREAMING_SNAKE_CASE_ :						Tuple    				= ids_tensor([self.batch_size, self.seq_length]			,     self.vocab_size							)

						SCREAMING_SNAKE_CASE_ :						List[Any]    				= None
						if self.use_input_mask:
										SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= random_attention_mask([self.batch_size, self.seq_length]							)

						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= GPTaConfig(
						    vocab_size=self.vocab_size			,     n_embd=self.hidden_size			,     n_layer=self.num_hidden_layers			,     n_head=self.num_attention_heads			,     n_positions=self.max_position_embeddings			,     use_cache=UpperCamelCase__			,     bos_token_id=self.bos_token_id			,     eos_token_id=self.eos_token_id			,     pad_token_id=self.pad_token_id			,     )

						return (config, input_ids, input_mask)


		def 	__snake_case      (    self							):
						SCREAMING_SNAKE_CASE_ :						Tuple    				= self.prepare_config_and_inputs()
						SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						List[str]    				= config_and_inputs
						SCREAMING_SNAKE_CASE_ :						Tuple    				= {'input_ids': input_ids, 'attention_mask': attention_mask}
						return config, inputs_dict


		def 	__snake_case      (    self							):
						SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Dict    				= self.prepare_config_and_inputs()

						SCREAMING_SNAKE_CASE_ :						int    				= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
						SCREAMING_SNAKE_CASE_ :						int    				= ids_tensor([self.batch_size, self.seq_length]			,     vocab_size=2							)

						return (
						    config,
						    input_ids,
						    attention_mask,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)


		def 	__snake_case      (    self			,     UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__							):
						SCREAMING_SNAKE_CASE_ :						Any    				= 20
						SCREAMING_SNAKE_CASE_ :						List[Any]    				= model_class_name(UpperCamelCase__							)

						SCREAMING_SNAKE_CASE_ :						str    				= model.init_cache(input_ids.shape[0]			,     UpperCamelCase__							)
						SCREAMING_SNAKE_CASE_ :						List[str]    				= jnp.ones((input_ids.shape[0], max_decoder_length)			,     dtype='i4'							)

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

						SCREAMING_SNAKE_CASE_ :						Tuple    				= jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,     dtype='i4'							)
						SCREAMING_SNAKE_CASE_ :						Tuple    				= model(
						    input_ids[:, -1:]			,     attention_mask=UpperCamelCase__			,     past_key_values=outputs_cache.past_key_values			,     position_ids=UpperCamelCase__			,     )

						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= model(UpperCamelCase__							)

						SCREAMING_SNAKE_CASE_ :						List[str]    				= 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 	__snake_case      (    self			,     UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__							):
						SCREAMING_SNAKE_CASE_ :						Tuple    				= 20
						SCREAMING_SNAKE_CASE_ :						str    				= model_class_name(UpperCamelCase__							)

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

						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= model.init_cache(input_ids.shape[0]			,     UpperCamelCase__							)
						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[int]    				= model(
						    input_ids[:, :-1]			,     attention_mask=UpperCamelCase__			,     past_key_values=UpperCamelCase__			,     position_ids=UpperCamelCase__			,     )
						SCREAMING_SNAKE_CASE_ :						str    				= jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,     dtype='i4'							)
						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= model(
						    input_ids[:, -1:]			,     past_key_values=outputs_cache.past_key_values			,     attention_mask=UpperCamelCase__			,     position_ids=UpperCamelCase__			,     )

						SCREAMING_SNAKE_CASE_ :						Tuple    				= model(UpperCamelCase__			,     attention_mask=UpperCamelCase__							)

						SCREAMING_SNAKE_CASE_ :						Union[str, 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 				__a				(	__A			,						__A			,						unittest.TestCase    ):





		'''simple docstring'''



		UpperCAmelCase__   :					Union[str, Any]					     =    (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
		UpperCAmelCase__   :					Optional[Any]					     =    (FlaxGPTaLMHeadModel,) if is_flax_available() else ()


		def 	__snake_case      (    self							):
						SCREAMING_SNAKE_CASE_ :						int    				= FlaxGPTaModelTester(self							)


		def 	__snake_case      (    self							):
						for model_class_name in self.all_model_classes:
										SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Tuple    				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward(UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__							)


		def 	__snake_case      (    self							):
						for model_class_name in self.all_model_classes:
										SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						str    				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward_with_attn_mask(
										    UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__			,     UpperCamelCase__							)


		@slow
		def 	__snake_case      (    self							):
						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= GPTaTokenizer.from_pretrained('gpt2'			,     pad_token='</s>'			,     padding_side='left'							)
						SCREAMING_SNAKE_CASE_ :						List[str]    				= tokenizer(['Hello this is a long string', 'Hey']			,     return_tensors='np'			,     padding=UpperCamelCase__			,     truncation=UpperCamelCase__							)

						SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= FlaxGPTaLMHeadModel.from_pretrained('gpt2'							)
						SCREAMING_SNAKE_CASE_ :						Any    				= False
						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= model.config.eos_token_id

						SCREAMING_SNAKE_CASE_ :						str    				= jax.jit(model.generate							)

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

						SCREAMING_SNAKE_CASE_ :						List[str]    				= tokenizer.batch_decode(UpperCamelCase__			,     skip_special_tokens=UpperCamelCase__							)

						SCREAMING_SNAKE_CASE_ :						List[str]    				= [
						    'Hello this is a long string of words. I\'m going to start with the first one.\n',
						    'Hey, I\'m not sure if I\'m going to be able to do',
						]

						self.assertListEqual(UpperCamelCase__			,     UpperCamelCase__							)


		@is_pt_flax_cross_test
		def 	__snake_case      (    self							):
						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:
										with self.subTest(model_class.__name__							):
														# prepare inputs
														SCREAMING_SNAKE_CASE_ :						List[Any]    				= self._prepare_for_class(UpperCamelCase__			,     UpperCamelCase__							)
														SCREAMING_SNAKE_CASE_ :						Any    				= {k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

														# load corresponding PyTorch class
														SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= model_class.__name__[4:]  # Skip the "Flax" at the beginning
														SCREAMING_SNAKE_CASE_ :						Optional[int]    				= getattr(UpperCamelCase__			,     UpperCamelCase__							)

														SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Tuple    				= pt_inputs['input_ids'].shape
														SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= np.random.randint(0			,     seq_length - 1			,     size=(batch_size,)							)
														for batch_idx, start_index in enumerate(UpperCamelCase__							):
																		SCREAMING_SNAKE_CASE_ :						str    				= 0
																		SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= 1
																		SCREAMING_SNAKE_CASE_ :						Optional[int]    				= 0
																		SCREAMING_SNAKE_CASE_ :						Dict    				= 1
														SCREAMING_SNAKE_CASE_ :						List[Any]    				= pt_model_class(UpperCamelCase__							).eval()
														SCREAMING_SNAKE_CASE_ :						Tuple    				= model_class(UpperCamelCase__			,     dtype=jnp.floataa							)

														SCREAMING_SNAKE_CASE_ :						Optional[int]    				= convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,     UpperCamelCase__							)
														SCREAMING_SNAKE_CASE_ :						int    				= fx_state

														with torch.no_grad():
																		SCREAMING_SNAKE_CASE_ :						int    				= pt_model(**UpperCamelCase__							).to_tuple()

														SCREAMING_SNAKE_CASE_ :						Tuple    				= fx_model(**UpperCamelCase__							).to_tuple()
														self.assertEqual(len(UpperCamelCase__							)			,     len(UpperCamelCase__							)			,     'Output lengths differ between Flax and PyTorch'							)
														for fx_output, pt_output in zip(UpperCamelCase__			,     UpperCamelCase__							):
																		self.assert_almost_equals(fx_output[:, -1]			,     pt_output[:, -1].numpy()			,     4E-2							)

														with tempfile.TemporaryDirectory() as tmpdirname:
																		pt_model.save_pretrained(UpperCamelCase__							)
																		SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= model_class.from_pretrained(UpperCamelCase__			,     from_pt=UpperCamelCase__							)

														SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= fx_model_loaded(**UpperCamelCase__							).to_tuple()
														self.assertEqual(
														    len(UpperCamelCase__							)			,     len(UpperCamelCase__							)			,     'Output lengths differ between Flax and PyTorch'							)
														for fx_output_loaded, pt_output in zip(UpperCamelCase__			,     UpperCamelCase__							):
																		self.assert_almost_equals(fx_output_loaded[:, -1]			,     pt_output[:, -1].numpy()			,     4E-2							)


		@is_pt_flax_cross_test
		def 	__snake_case      (    self							):
						SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Any    				= self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__							):
														# prepare inputs
														SCREAMING_SNAKE_CASE_ :						Any    				= self._prepare_for_class(UpperCamelCase__			,     UpperCamelCase__							)
														SCREAMING_SNAKE_CASE_ :						Dict    				= {k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

														# load corresponding PyTorch class
														SCREAMING_SNAKE_CASE_ :						str    				= model_class.__name__[4:]  # Skip the "Flax" at the beginning
														SCREAMING_SNAKE_CASE_ :						Any    				= getattr(UpperCamelCase__			,     UpperCamelCase__							)

														SCREAMING_SNAKE_CASE_ :						str    				= pt_model_class(UpperCamelCase__							).eval()
														SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= model_class(UpperCamelCase__			,     dtype=jnp.floataa							)

														SCREAMING_SNAKE_CASE_ :						str    				= load_flax_weights_in_pytorch_model(UpperCamelCase__			,     fx_model.params							)
														SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Dict    				= pt_inputs['input_ids'].shape
														SCREAMING_SNAKE_CASE_ :						List[Any]    				= np.random.randint(0			,     seq_length - 1			,     size=(batch_size,)							)
														for batch_idx, start_index in enumerate(UpperCamelCase__							):
																		SCREAMING_SNAKE_CASE_ :						Optional[int]    				= 0
																		SCREAMING_SNAKE_CASE_ :						Tuple    				= 1
																		SCREAMING_SNAKE_CASE_ :						List[Any]    				= 0
																		SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= 1

														# make sure weights are tied in PyTorch
														pt_model.tie_weights()

														with torch.no_grad():
																		SCREAMING_SNAKE_CASE_ :						List[Any]    				= pt_model(**UpperCamelCase__							).to_tuple()

														SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= fx_model(**UpperCamelCase__							).to_tuple()
														self.assertEqual(len(UpperCamelCase__							)			,     len(UpperCamelCase__							)			,     'Output lengths differ between Flax and PyTorch'							)
														for fx_output, pt_output in zip(UpperCamelCase__			,     UpperCamelCase__							):
																		self.assert_almost_equals(fx_output[:, -1]			,     pt_output[:, -1].numpy()			,     4E-2							)

														with tempfile.TemporaryDirectory() as tmpdirname:
																		fx_model.save_pretrained(UpperCamelCase__							)
																		SCREAMING_SNAKE_CASE_ :						Optional[int]    				= pt_model_class.from_pretrained(UpperCamelCase__			,     from_flax=UpperCamelCase__							)

														with torch.no_grad():
																		SCREAMING_SNAKE_CASE_ :						Tuple    				= pt_model_loaded(**UpperCamelCase__							).to_tuple()

														self.assertEqual(
														    len(UpperCamelCase__							)			,     len(UpperCamelCase__							)			,     'Output lengths differ between Flax and PyTorch'							)
														for fx_output, pt_output in zip(UpperCamelCase__			,     UpperCamelCase__							):
																		self.assert_almost_equals(fx_output[:, -1]			,     pt_output[:, -1].numpy()			,     4E-2							)






		@slow
		def 	__snake_case      (    self							):
						for model_class_name in self.all_model_classes:
										SCREAMING_SNAKE_CASE_ :						int    				= model_class_name.from_pretrained('gpt2'			,     from_pt=UpperCamelCase__							)
										SCREAMING_SNAKE_CASE_ :						List[Any]    				= model(np.ones((1, 1)							)							)
										self.assertIsNotNone(UpperCamelCase__							)







from ...configuration_utils import PretrainedConfig
from ...utils import logging


A										=      logging.get_logger(__name__)

A										=      {
    'ZinengTang/tvlt-base': 'https://huggingface.co/ZinengTang/tvlt-base/blob/main/config.json',
}


class 				__a				(	__A    ):





		'''simple docstring'''



		UpperCAmelCase__   :					str					     =    """tvlt"""


		def __init__(    self			,     UpperCamelCase__=224			,     UpperCamelCase__=2048			,     UpperCamelCase__=128			,     UpperCamelCase__=[16, 16]			,     UpperCamelCase__=[16, 16]			,     UpperCamelCase__=3			,     UpperCamelCase__=1			,     UpperCamelCase__=8			,     UpperCamelCase__=768			,     UpperCamelCase__=12			,     UpperCamelCase__=12			,     UpperCamelCase__=3072			,     UpperCamelCase__="gelu"			,     UpperCamelCase__=0.0			,     UpperCamelCase__=0.0			,     UpperCamelCase__=0.02			,     UpperCamelCase__=1E-6			,     UpperCamelCase__=True			,     UpperCamelCase__=False			,     UpperCamelCase__=16			,     UpperCamelCase__=512			,     UpperCamelCase__=8			,     UpperCamelCase__=2048			,     UpperCamelCase__=0.75			,     UpperCamelCase__=0.15			,     UpperCamelCase__="frame-level"			,     UpperCamelCase__=True			,     UpperCamelCase__=True			,     UpperCamelCase__="classification"			,     **UpperCamelCase__			,     ):
						super().__init__(**UpperCamelCase__							)

						if audio_mask_type not in ("frame-level", "patch_level"):
										raise ValueError(
										    'audio_mask_type must be one of two acceptable strategies - {\'frame_level\', \'patch-level\') '
										    F'''got {audio_mask_type}'''							)

						SCREAMING_SNAKE_CASE_ :						Tuple    				= image_size
						SCREAMING_SNAKE_CASE_ :						int    				= spectrogram_length
						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= frequency_length
						SCREAMING_SNAKE_CASE_ :						Any    				= image_patch_size
						SCREAMING_SNAKE_CASE_ :						List[Any]    				= audio_patch_size
						SCREAMING_SNAKE_CASE_ :						int    				= num_image_channels
						SCREAMING_SNAKE_CASE_ :						Any    				= num_audio_channels
						SCREAMING_SNAKE_CASE_ :						Any    				= num_frames

						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= hidden_size
						SCREAMING_SNAKE_CASE_ :						str    				= num_hidden_layers
						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= num_attention_heads
						SCREAMING_SNAKE_CASE_ :						str    				= intermediate_size
						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= hidden_act
						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= hidden_dropout_prob
						SCREAMING_SNAKE_CASE_ :						int    				= attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE_ :						Dict    				= initializer_range
						SCREAMING_SNAKE_CASE_ :						Dict    				= layer_norm_eps
						SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= qkv_bias
						SCREAMING_SNAKE_CASE_ :						List[str]    				= use_mean_pooling

						SCREAMING_SNAKE_CASE_ :						str    				= decoder_num_attention_heads
						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= decoder_hidden_size
						SCREAMING_SNAKE_CASE_ :						Any    				= decoder_num_hidden_layers
						SCREAMING_SNAKE_CASE_ :						List[Any]    				= decoder_intermediate_size
						SCREAMING_SNAKE_CASE_ :						List[str]    				= pixel_mask_ratio
						SCREAMING_SNAKE_CASE_ :						Optional[int]    				= audio_mask_ratio
						SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= audio_mask_type

						SCREAMING_SNAKE_CASE_ :						Any    				= task_matching
						SCREAMING_SNAKE_CASE_ :						Any    				= task_mae
						SCREAMING_SNAKE_CASE_ :						int    				= loss_type