





import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    FunnelBaseModel,
	    FunnelForMaskedLM,
	    FunnelForMultipleChoice,
	    FunnelForPreTraining,
	    FunnelForQuestionAnswering,
	    FunnelForSequenceClassification,
	    FunnelForTokenClassification,
	    FunnelModel,
	)



class       __lowerCAmelCase :






			def __init__(self      ,							__magic_name__      ,							__magic_name__=13      ,							__magic_name__=7      ,							__magic_name__=True      ,							__magic_name__=True      ,							__magic_name__=True      ,							__magic_name__=True      ,							__magic_name__=99      ,							__magic_name__=[1, 1, 2]      ,							__magic_name__=1      ,							__magic_name__=32      ,							__magic_name__=4      ,							__magic_name__=8      ,							__magic_name__=37      ,							__magic_name__="gelu_new"      ,							__magic_name__=0.1      ,							__magic_name__=0.1      ,							__magic_name__=0.0      ,							__magic_name__=512      ,							__magic_name__=3      ,							__magic_name__=0.02      ,							__magic_name__=3      ,							__magic_name__=4      ,							__magic_name__=None      ,							__magic_name__=False      ,							)  ->     Any:


							'''simple docstring'''



							snake_case_		:					str			    =						parent
							snake_case_		:					Optional[Any]			    =						batch_size
							snake_case_		:					Any			    =						seq_length
							snake_case_		:					str			    =						is_training
							snake_case_		:					int			    =						use_input_mask
							snake_case_		:					List[str]			    =						use_token_type_ids
							snake_case_		:					Optional[int]			    =						use_labels
							snake_case_		:					Any			    =						vocab_size
							snake_case_		:					Any			    =						block_sizes
							snake_case_		:					Optional[Any]			    =						num_decoder_layers
							snake_case_		:					Optional[int]			    =						d_model
							snake_case_		:					Optional[Any]			    =						n_head
							snake_case_		:					List[str]			    =						d_head
							snake_case_		:					Tuple			    =						d_inner
							snake_case_		:					Dict			    =						hidden_act
							snake_case_		:					str			    =						hidden_dropout
							snake_case_		:					List[str]			    =						attention_dropout
							snake_case_		:					Optional[int]			    =						activation_dropout
							snake_case_		:					Optional[Any]			    =						max_position_embeddings
							snake_case_		:					List[str]			    =						type_vocab_size
							snake_case_		:					Tuple			    =						2
							snake_case_		:					Tuple			    =						num_labels
							snake_case_		:					int			    =						num_choices
							snake_case_		:					Dict			    =						scope
							snake_case_		:					Tuple			    =						initializer_std

							# Used in the tests to check the size of the first attention layer
							snake_case_		:					List[str]			    =						n_head
							# Used in the tests to check the size of the first hidden state
							snake_case_		:					Any			    =						self.d_model
							# Used in the tests to check the number of output hidden states/attentions
							snake_case_		:					List[str]			    =						sum(self.block_sizes					) + (0 if base else self.num_decoder_layers)
							# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
							# the last hidden state of the first block (which is the first hidden state of the decoder).
							if not base:
											snake_case_		:					str			    =						self.num_hidden_layers + 2






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


							'''simple docstring'''



							snake_case_		:					Any			    =						ids_tensor([self.batch_size, self.seq_length]      ,							self.vocab_size					)

							snake_case_		:					int			    =						None
							if self.use_input_mask:
											snake_case_		:					Any			    =						ids_tensor([self.batch_size, self.seq_length]      ,							vocab_size=2					)

							snake_case_		:					List[str]			    =						None
							if self.use_token_type_ids:
											snake_case_		:					Optional[int]			    =						ids_tensor([self.batch_size, self.seq_length]      ,							self.type_vocab_size					)

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

							snake_case_		:					int			    =						self.get_config()

							return (
							    config,
							    input_ids,
							    token_type_ids,
							    input_mask,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							    fake_token_labels,
							)






			def     lowerCamelCase						(self					)  ->     Any:


							'''simple docstring'''



							return FunnelConfig(
							    vocab_size=self.vocab_size      ,							block_sizes=self.block_sizes      ,							num_decoder_layers=self.num_decoder_layers      ,							d_model=self.d_model      ,							n_head=self.n_head      ,							d_head=self.d_head      ,							d_inner=self.d_inner      ,							hidden_act=self.hidden_act      ,							hidden_dropout=self.hidden_dropout      ,							attention_dropout=self.attention_dropout      ,							activation_dropout=self.activation_dropout      ,							max_position_embeddings=self.max_position_embeddings      ,							type_vocab_size=self.type_vocab_size      ,							initializer_std=self.initializer_std      ,							)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     Optional[Any]:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						FunnelModel(config=__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					str			    =						model(__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__					)
							snake_case_		:					List[str]			    =						model(__magic_name__      ,							token_type_ids=__magic_name__					)
							snake_case_		:					Tuple			    =						model(__magic_name__					)
							self.parent.assertEqual(result.last_hidden_state.shape      ,							(self.batch_size, self.seq_length, self.d_model)					)

							snake_case_		:					Optional[int]			    =						False
							snake_case_		:					Union[str, Any]			    =						model(__magic_name__					)
							self.parent.assertEqual(result.last_hidden_state.shape      ,							(self.batch_size, self.seq_length, self.d_model)					)

							snake_case_		:					List[str]			    =						False
							snake_case_		:					int			    =						model(__magic_name__					)
							self.parent.assertEqual(result.last_hidden_state.shape      ,							(self.batch_size, self.seq_length, self.d_model)					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     str:


							'''simple docstring'''



							snake_case_		:					Optional[Any]			    =						FunnelBaseModel(config=__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					Union[str, Any]			    =						model(__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__					)
							snake_case_		:					List[Any]			    =						model(__magic_name__      ,							token_type_ids=__magic_name__					)
							snake_case_		:					List[str]			    =						model(__magic_name__					)
							self.parent.assertEqual(result.last_hidden_state.shape      ,							(self.batch_size, 2, self.d_model)					)

							snake_case_		:					List[str]			    =						False
							snake_case_		:					Tuple			    =						model(__magic_name__					)
							self.parent.assertEqual(result.last_hidden_state.shape      ,							(self.batch_size, 3, self.d_model)					)

							snake_case_		:					Tuple			    =						False
							snake_case_		:					Optional[Any]			    =						model(__magic_name__					)
							self.parent.assertEqual(result.last_hidden_state.shape      ,							(self.batch_size, 2, self.d_model)					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     Optional[Any]:


							'''simple docstring'''



							snake_case_		:					Tuple			    =						self.num_labels
							snake_case_		:					Dict			    =						FunnelForPreTraining(config=__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					Any			    =						model(__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							labels=__magic_name__					)
							self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.seq_length)					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     Optional[int]:


							'''simple docstring'''



							snake_case_		:					Dict			    =						FunnelForMaskedLM(config=__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					Tuple			    =						model(__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							labels=__magic_name__					)
							self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.seq_length, self.vocab_size)					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     int:


							'''simple docstring'''



							snake_case_		:					Union[str, Any]			    =						self.num_labels
							snake_case_		:					Union[str, Any]			    =						FunnelForSequenceClassification(__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					Any			    =						model(__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							labels=__magic_name__					)
							self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.num_labels)					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     Optional[int]:


							'''simple docstring'''



							snake_case_		:					Tuple			    =						self.num_choices
							snake_case_		:					Optional[Any]			    =						FunnelForMultipleChoice(config=__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					Union[str, Any]			    =						input_ids.unsqueeze(1					).expand(-1      ,							self.num_choices      ,							-1					).contiguous()
							snake_case_		:					List[str]			    =						token_type_ids.unsqueeze(1					).expand(-1      ,							self.num_choices      ,							-1					).contiguous()
							snake_case_		:					int			    =						input_mask.unsqueeze(1					).expand(-1      ,							self.num_choices      ,							-1					).contiguous()
							snake_case_		:					List[str]			    =						model(
							    __magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							labels=__magic_name__      ,							)
							self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.num_choices)					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     Optional[int]:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						self.num_labels
							snake_case_		:					Optional[Any]			    =						FunnelForTokenClassification(config=__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					Tuple			    =						model(__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							labels=__magic_name__					)
							self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.seq_length, self.num_labels)					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							)  ->     Tuple:


							'''simple docstring'''



							snake_case_		:					int			    =						FunnelForQuestionAnswering(config=__magic_name__					)
							model.to(__magic_name__					)
							model.eval()
							snake_case_		:					List[Any]			    =						model(
							    __magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							start_positions=__magic_name__      ,							end_positions=__magic_name__      ,							)
							self.parent.assertEqual(result.start_logits.shape      ,							(self.batch_size, self.seq_length)					)
							self.parent.assertEqual(result.end_logits.shape      ,							(self.batch_size, self.seq_length)					)






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


							'''simple docstring'''



							snake_case_		:					List[Any]			    =						self.prepare_config_and_inputs()
							(
							    (
							    snake_case_
							)    ,							(
							    snake_case_
							)    ,							(
							    snake_case_
							)    ,							(
							    snake_case_
							)    ,							(
							    snake_case_
							)    ,							(
							    snake_case_
							)    ,							(
							    snake_case_
							)    ,							(
							    snake_case_
							)    ,							
							)		:					int			    =						config_and_inputs
							snake_case_		:					int			    =						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
							return config, inputs_dict




@require_torch
class       __lowerCAmelCase (    _a,  _a,  unittest.TestCase  ):
			lowerCamelCase_  :  List[Any]     =   False
			lowerCamelCase_  :  str     =   False
			lowerCamelCase_  :  Any     =   (
			    (
			        FunnelModel,
			        FunnelForMaskedLM,
			        FunnelForPreTraining,
			        FunnelForQuestionAnswering,
			        FunnelForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase_  :  Tuple     =   (
			    {
			        '''feature-extraction''': (FunnelBaseModel, FunnelModel),
			        '''fill-mask''': FunnelForMaskedLM,
			        '''question-answering''': FunnelForQuestionAnswering,
			        '''text-classification''': FunnelForSequenceClassification,
			        '''token-classification''': FunnelForTokenClassification,
			        '''zero-shot''': FunnelForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__=False					)  ->     Any:


							'''simple docstring'''



							snake_case_		:					List[Any]			    =						super()._prepare_for_class(__magic_name__      ,							__magic_name__      ,							return_labels=__magic_name__					)

							if return_labels:
											if model_class in get_values(__magic_name__					):
															snake_case_		:					Tuple			    =						torch.zeros(
															    (self.model_tester.batch_size, self.model_tester.seq_length)      ,							dtype=torch.long      ,							device=__magic_name__					)
							return inputs_dict






			def     lowerCamelCase						(self					)  ->     Optional[int]:


							'''simple docstring'''



							snake_case_		:					Union[str, Any]			    =						FunnelModelTester(self					)
							snake_case_		:					List[Any]			    =						ConfigTester(self      ,							config_class=__magic_name__					)






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


							'''simple docstring'''



							self.config_tester.run_common_tests()






			def     lowerCamelCase						(self					)  ->     Optional[int]:


							'''simple docstring'''



							snake_case_		:					str			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__magic_name__					)






			def     lowerCamelCase						(self					)  ->     Dict:


							'''simple docstring'''



							snake_case_		:					List[Any]			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*__magic_name__					)






			def     lowerCamelCase						(self					)  ->     str:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__magic_name__					)






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


							'''simple docstring'''



							snake_case_		:					int			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__magic_name__					)






			def     lowerCamelCase						(self					)  ->     Optional[Any]:


							'''simple docstring'''



							snake_case_		:					List[str]			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__magic_name__					)






			def     lowerCamelCase						(self      ,							__magic_name__					)  ->     Any:


							'''simple docstring'''



							if hasattr(__magic_name__      ,							'''weight'''					) and module.weight is not None:
											module.weight.data.fill_(3					)
							if hasattr(__magic_name__      ,							'''bias'''					) and module.bias is not None:
											module.bias.data.fill_(3					)

							for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
											if hasattr(__magic_name__      ,							__magic_name__					) and getattr(__magic_name__      ,							__magic_name__					) is not None:
															snake_case_		:					Tuple			    =						getattr(__magic_name__      ,							__magic_name__					)
															weight.data.fill_(3					)




@require_torch
class       __lowerCAmelCase (    _a,  unittest.TestCase  ):
			lowerCamelCase_  :  int     =   False
			lowerCamelCase_  :  Any     =   False
			lowerCamelCase_  :  int     =   (
			    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
			)






			def     lowerCamelCase						(self					)  ->     int:


							'''simple docstring'''



							snake_case_		:					Dict			    =						FunnelModelTester(self      ,							base=__magic_name__					)
							snake_case_		:					Tuple			    =						ConfigTester(self      ,							config_class=__magic_name__					)






			def     lowerCamelCase						(self					)  ->     Dict:


							'''simple docstring'''



							self.config_tester.run_common_tests()






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


							'''simple docstring'''



							snake_case_		:					Dict			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_base_model(*__magic_name__					)






			def     lowerCamelCase						(self					)  ->     Tuple:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__magic_name__					)






			def     lowerCamelCase						(self					)  ->     str:


							'''simple docstring'''



							snake_case_		:					Dict			    =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__magic_name__					)






			def     lowerCamelCase						(self					)  ->     int:


							'''simple docstring'''



							snake_case_    ,							snake_case_		:					Any			    =						self.model_tester.prepare_config_and_inputs_for_common()
							snake_case_		:					List[str]			    =						True

							for model_class in self.all_model_classes:
											if model_class.__name__ == "FunnelBaseModel":
															continue
											snake_case_		:					Optional[int]			    =						model_class(__magic_name__					)
											model.to(__magic_name__					)
											model.train()
											snake_case_		:					Tuple			    =						self._prepare_for_class(__magic_name__      ,							__magic_name__      ,							return_labels=__magic_name__					)
											snake_case_		:					Any			    =						model(**__magic_name__					).loss
											loss.backward()






			def     lowerCamelCase						(self      ,							__magic_name__					)  ->     Optional[int]:


							'''simple docstring'''



							if hasattr(__magic_name__      ,							'''weight'''					) and module.weight is not None:
											module.weight.data.fill_(3					)
							if hasattr(__magic_name__      ,							'''bias'''					) and module.bias is not None:
											module.bias.data.fill_(3					)

							for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
											if hasattr(__magic_name__      ,							__magic_name__					) and getattr(__magic_name__      ,							__magic_name__					) is not None:
															snake_case_		:					Any			    =						getattr(__magic_name__      ,							__magic_name__					)
															weight.data.fill_(3					)




@require_torch
@require_sentencepiece
@require_tokenizers
class       __lowerCAmelCase (    unittest.TestCase  ):






			def     lowerCamelCase						(self					)  ->     Optional[Any]:


							'''simple docstring'''



							snake_case_		:					Tuple			    =						13
							snake_case_		:					Tuple			    =						7
							snake_case_		:					List[str]			    =						torch.arange(0      ,							batch_size * sequence_length					).long().reshape(__magic_name__      ,							__magic_name__					)
							snake_case_		:					List[Any]			    =						[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
							snake_case_		:					Optional[Any]			    =						torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]					)

							snake_case_		:					Union[str, Any]			    =						FunnelModel.from_pretrained('''sgugger/funnel-random-tiny'''					)
							snake_case_		:					Optional[Any]			    =						model(__magic_name__      ,							token_type_ids=__magic_name__					)[0].abs()

							snake_case_		:					Optional[Any]			    =						torch.tensor(2_344.8_352					)
							snake_case_		:					Optional[int]			    =						torch.tensor(0.8_052					)
							self.assertTrue(torch.allclose(output.sum()      ,							__magic_name__      ,							atol=1e-4					)					)
							self.assertTrue(torch.allclose(output.mean()      ,							__magic_name__      ,							atol=1e-4					)					)

							snake_case_		:					List[Any]			    =						torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]					)
							snake_case_		:					Dict			    =						model(__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__					)[0].abs()

							snake_case_		:					Union[str, Any]			    =						torch.tensor(2_343.8_425					)
							snake_case_		:					List[Any]			    =						torch.tensor(0.8_049					)
							self.assertTrue(torch.allclose(output.sum()      ,							__magic_name__      ,							atol=1e-4					)					)
							self.assertTrue(torch.allclose(output.mean()      ,							__magic_name__      ,							atol=1e-4					)					)






			@slow
			def     lowerCamelCase						(self					)  ->     Optional[Any]:


							'''simple docstring'''



							snake_case_		:					Any			    =						FunnelTokenizer.from_pretrained('''huggingface/funnel-small'''					)
							snake_case_		:					Optional[int]			    =						FunnelModel.from_pretrained('''huggingface/funnel-small'''					)
							snake_case_		:					Union[str, Any]			    =						tokenizer('''Hello! I am the Funnel Transformer model.'''      ,							return_tensors='''pt'''					)
							snake_case_		:					Tuple			    =						model(**__magic_name__					)[0]

							snake_case_		:					List[Any]			    =						torch.tensor(235.7_246					)
							snake_case_		:					Tuple			    =						torch.tensor(0.0_256					)
							self.assertTrue(torch.allclose(output.sum()      ,							__magic_name__      ,							atol=1e-4					)					)
							self.assertTrue(torch.allclose(output.mean()      ,							__magic_name__      ,							atol=1e-4					)					)







from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


lowerCAmelCase_												=       logging.get_logger(__name__)



class       __lowerCAmelCase (    nn.Module  ):
			lowerCamelCase_  :  HybridCLIPConfig
			lowerCamelCase_  :  jnp.dtype     =   jnp.floataa






			def     lowerCamelCase						(self					)  ->     int:


							'''simple docstring'''



							snake_case_		:					Any			    =						self.config.text_config
							snake_case_		:					Any			    =						self.config.vision_config

							snake_case_		:					Union[str, Any]			    =						self.config.projection_dim
							snake_case_		:					Dict			    =						text_config.hidden_size
							snake_case_		:					Union[str, Any]			    =						vision_config.hidden_size

							snake_case_		:					Union[str, Any]			    =						FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
							snake_case_		:					Dict			    =						FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__      ,							__magic_name__					).module_class

							snake_case_		:					Optional[Any]			    =						text_module(__magic_name__      ,							dtype=self.dtype					)
							snake_case_		:					Dict			    =						vision_module(__magic_name__      ,							dtype=self.dtype					)

							snake_case_		:					Dict			    =						nn.Dense(
							    self.projection_dim      ,							dtype=self.dtype      ,							kernel_init=jax.nn.initializers.normal(0.02					)      ,							use_bias=__magic_name__      ,							)
							snake_case_		:					Optional[Any]			    =						nn.Dense(
							    self.projection_dim      ,							dtype=self.dtype      ,							kernel_init=jax.nn.initializers.normal(0.02					)      ,							use_bias=__magic_name__      ,							)
							snake_case_		:					Optional[Any]			    =						self.param('''logit_scale'''      ,							jax.nn.initializers.ones      ,							[]					)






			def __call__(self      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = True      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							)  ->     Any:


							'''simple docstring'''



							snake_case_		:					Tuple			    =						return_dict if return_dict is not None else self.config.return_dict

							snake_case_		:					Tuple			    =						self.vision_model(
							    pixel_values=__magic_name__      ,							deterministic=__magic_name__      ,							output_attentions=__magic_name__      ,							output_hidden_states=__magic_name__      ,							return_dict=__magic_name__      ,							)

							snake_case_		:					Tuple			    =						self.text_model(
							    input_ids=__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							position_ids=__magic_name__      ,							deterministic=__magic_name__      ,							output_attentions=__magic_name__      ,							output_hidden_states=__magic_name__      ,							return_dict=__magic_name__      ,							)

							snake_case_		:					Dict			    =						vision_outputs[1]
							snake_case_		:					Optional[Any]			    =						self.visual_projection(__magic_name__					)

							snake_case_		:					List[str]			    =						text_outputs[1]
							snake_case_		:					Union[str, Any]			    =						self.text_projection(__magic_name__					)

							# normalized features
							snake_case_		:					List[str]			    =						image_embeds / jnp.linalg.norm(__magic_name__      ,							axis=-1      ,							keepdims=__magic_name__					)
							snake_case_		:					Any			    =						text_embeds / jnp.linalg.norm(__magic_name__      ,							axis=-1      ,							keepdims=__magic_name__					)

							# cosine similarity as logits
							snake_case_		:					Optional[Any]			    =						jnp.exp(self.logit_scale					)
							snake_case_		:					str			    =						jnp.matmul(__magic_name__      ,							image_embeds.T					) * logit_scale
							snake_case_		:					List[str]			    =						logits_per_text.T

							if not return_dict:
											return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

							return FlaxCLIPOutput(
							    logits_per_image=__magic_name__      ,							logits_per_text=__magic_name__      ,							text_embeds=__magic_name__      ,							image_embeds=__magic_name__      ,							text_model_output=__magic_name__      ,							vision_model_output=__magic_name__      ,							)




class       __lowerCAmelCase (    _a  ):
			lowerCamelCase_  :  Dict     =   HybridCLIPConfig
			lowerCamelCase_  :  int     =   FlaxHybridCLIPModule






			def __init__(self      ,							__magic_name__      ,							__magic_name__ = None      ,							__magic_name__ = 0      ,							__magic_name__ = jnp.floataa      ,							**__magic_name__      ,							)  ->     str:


							'''simple docstring'''



							if input_shape is None:
											snake_case_		:					int			    =						((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

							snake_case_		:					Union[str, Any]			    =						self.module_class(config=__magic_name__      ,							dtype=__magic_name__      ,							**__magic_name__					)
							super().__init__(__magic_name__      ,							__magic_name__      ,							input_shape=__magic_name__      ,							seed=__magic_name__      ,							dtype=__magic_name__					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__ = None					)  ->     FrozenDict:


							'''simple docstring'''



							snake_case_		:					List[Any]			    =						jnp.zeros(input_shape[0]      ,							dtype='''i4'''					)
							snake_case_		:					Optional[Any]			    =						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__					).shape[-1]					)      ,							input_shape[0]					)
							snake_case_		:					Union[str, Any]			    =						jnp.ones_like(__magic_name__					)
							snake_case_		:					List[Any]			    =						jnp.ones_like(__magic_name__					)

							snake_case_		:					Union[str, Any]			    =						jax.random.normal(__magic_name__      ,							input_shape[1]					)

							snake_case_    ,							snake_case_		:					Optional[int]			    =						jax.random.split(__magic_name__					)
							snake_case_		:					int			    =						{'''params''': params_rng, '''dropout''': dropout_rng}

							return self.module.init(__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__					)["params"]






			def __call__(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__ = False      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							)  ->     Any:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						output_attentions if output_attentions is not None else self.config.output_attentions
							snake_case_		:					Tuple			    =						(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							snake_case_		:					List[Any]			    =						return_dict if return_dict is not None else self.config.return_dict

							if position_ids is None:
											snake_case_		:					Tuple			    =						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__					).shape[-1]					)      ,							input_ids.shape					)

							if token_type_ids is None:
											snake_case_		:					Tuple			    =						jnp.zeros_like(__magic_name__					)

							if attention_mask is None:
											snake_case_		:					Optional[Any]			    =						jnp.ones_like(__magic_name__					)

							# Handle any PRNG if needed
							snake_case_		:					List[str]			    =						{}
							if dropout_rng is not None:
											snake_case_		:					Tuple			    =						dropout_rng

							return self.module.apply(
							    {'''params''': params or self.params}      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype=jnp.floataa					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							not train      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							rngs=__magic_name__      ,							)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__=False      ,							)  ->     Union[str, Any]:


							'''simple docstring'''



							if position_ids is None:
											snake_case_		:					Tuple			    =						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__					).shape[-1]					)      ,							input_ids.shape					)

							if token_type_ids is None:
											snake_case_		:					List[str]			    =						jnp.zeros_like(__magic_name__					)

							if attention_mask is None:
											snake_case_		:					Union[str, Any]			    =						jnp.ones_like(__magic_name__					)

							# Handle any PRNG if needed
							snake_case_		:					Optional[int]			    =						{}
							if dropout_rng is not None:
											snake_case_		:					int			    =						dropout_rng

							def _get_features(__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__					):
											snake_case_		:					int			    =						module.text_model(
											    input_ids=__magic_name__      ,							attention_mask=__magic_name__      ,							position_ids=__magic_name__      ,							token_type_ids=__magic_name__      ,							deterministic=__magic_name__      ,							)
											snake_case_		:					Optional[int]			    =						text_outputs[1]
											snake_case_		:					int			    =						module.text_projection(__magic_name__					)
											return text_features

							return self.module.apply(
							    {'''params''': params or self.params}      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							not train      ,							method=_get_features      ,							rngs=__magic_name__      ,							)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__=False					)  ->     Dict:


							'''simple docstring'''



							snake_case_		:					Any			    =						{}
							if dropout_rng is not None:
											snake_case_		:					Optional[int]			    =						dropout_rng

							def _get_features(__magic_name__      ,							__magic_name__      ,							__magic_name__					):
											snake_case_		:					Dict			    =						module.vision_model(pixel_values=__magic_name__      ,							deterministic=__magic_name__					)
											snake_case_		:					str			    =						vision_outputs[1]  # pooled_output
											snake_case_		:					str			    =						module.visual_projection(__magic_name__					)
											return image_features

							return self.module.apply(
							    {'''params''': params or self.params}      ,							jnp.array(__magic_name__      ,							dtype=jnp.floataa					)      ,							not train      ,							method=_get_features      ,							rngs=__magic_name__      ,							)






			@classmethod
			def     lowerCamelCase						(cls      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							*__magic_name__      ,							**__magic_name__      ,							)  ->     FlaxPreTrainedModel:


							'''simple docstring'''



							snake_case_		:					Optional[Any]			    =						{
							    argument[len('''text_'''					) :]: value for argument, value in kwargs.items() if argument.startswith('''text_'''					)
							}

							snake_case_		:					Dict			    =						{
							    argument[len('''vision_'''					) :]: value for argument, value in kwargs.items() if argument.startswith('''vision_'''					)
							}

							# remove text, vision kwargs from kwargs
							for key in kwargs_text.keys():
											del kwargs["text_" + key]
							for key in kwargs_vision.keys():
											del kwargs["vision_" + key]

							# Load and initialize the text and vision model
							snake_case_		:					str			    =						kwargs_text.pop('''model'''      ,							__magic_name__					)
							if text_model is None:
											assert (
											    text_model_name_or_path is not None
											), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
											from transformers import FlaxAutoModel

											if "config" not in kwargs_text:
															from transformers import AutoConfig

															snake_case_		:					Dict			    =						AutoConfig.from_pretrained(__magic_name__					)
															snake_case_		:					int			    =						text_config

											snake_case_		:					Optional[int]			    =						FlaxAutoModel.from_pretrained(__magic_name__      ,							*__magic_name__      ,							**__magic_name__					)

							snake_case_		:					Any			    =						kwargs_vision.pop('''model'''      ,							__magic_name__					)
							if vision_model is None:
											assert (
											    vision_model_name_or_path is not None
											), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
											from transformers import FlaxAutoModel

											if "config" not in kwargs_vision:
															from transformers import AutoConfig

															snake_case_		:					Optional[int]			    =						AutoConfig.from_pretrained(__magic_name__					)
															snake_case_		:					List[str]			    =						vision_config

											snake_case_		:					List[Any]			    =						FlaxAutoModel.from_pretrained(__magic_name__      ,							*__magic_name__      ,							**__magic_name__					)

							# instantiate config with corresponding kwargs
							snake_case_		:					Dict			    =						kwargs.pop('''dtype'''      ,							jnp.floataa					)
							snake_case_		:					str			    =						HybridCLIPConfig.from_text_vision_configs(text_model.config      ,							vision_model.config      ,							**__magic_name__					)

							# init model
							snake_case_		:					int			    =						cls(__magic_name__      ,							*__magic_name__      ,							dtype=__magic_name__      ,							**__magic_name__					)

							if vision_config.model_type == "clip":
											snake_case_		:					int			    =						vision_model.params['''vision_model''']
											snake_case_		:					Any			    =						vision_model.params['''visual_projection''']['''kernel''']
							else:
											snake_case_		:					List[Any]			    =						vision_model.params

							snake_case_		:					Union[str, Any]			    =						text_model.params

							return model
