






from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

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


if is_tf_available():
	import tensorflow as tf

	from transformers import (
	    TFFunnelBaseModel,
	    TFFunnelForMaskedLM,
	    TFFunnelForMultipleChoice,
	    TFFunnelForPreTraining,
	    TFFunnelForQuestionAnswering,
	    TFFunnelForSequenceClassification,
	    TFFunnelForTokenClassification,
	    TFFunnelModel,
	)



class 							SCREAMING_SNAKE_CASE__		:





					"""simple docstring"""






					def __init__(      self  :						Union[str, Any]							,	__A  :						int							,	__A  :						List[Any]=1_3							,	__A  :						str=7							,	__A  :						Union[str, Any]=True							,	__A  :						Tuple=True							,	__A  :						Tuple=True							,	__A  :						Optional[Any]=True							,	__A  :						List[Any]=9_9							,	__A  :						List[str]=[1, 1, 2]							,	__A  :						Optional[Any]=1							,	__A  :						int=3_2							,	__A  :						Optional[Any]=4							,	__A  :						Optional[Any]=8							,	__A  :						List[str]=3_7							,	__A  :						Optional[int]="gelu_new"							,	__A  :						str=0.1							,	__A  :						Union[str, Any]=0.1							,	__A  :						Union[str, Any]=0.0							,	__A  :						List[Any]=5_1_2							,	__A  :						List[Any]=3							,	__A  :						Optional[Any]=0.0_2							,	__A  :						Optional[int]=3							,	__A  :						List[str]=4							,	__A  :						int=None							,	__A  :						Tuple=False							,	):
						snake_case__				:   List[Any]     =						parent
						snake_case__				:   Optional[int]     =						batch_size
						snake_case__				:   Dict     =						seq_length
						snake_case__				:   Tuple     =						is_training
						snake_case__				:   Dict     =						use_input_mask
						snake_case__				:   Union[str, Any]     =						use_token_type_ids
						snake_case__				:   Any     =						use_labels
						snake_case__				:   Union[str, Any]     =						vocab_size
						snake_case__				:   Any     =						block_sizes
						snake_case__				:   Optional[Any]     =						num_decoder_layers
						snake_case__				:   Dict     =						d_model
						snake_case__				:   Dict     =						n_head
						snake_case__				:   Union[str, Any]     =						d_head
						snake_case__				:   Dict     =						d_inner
						snake_case__				:   List[str]     =						hidden_act
						snake_case__				:   int     =						hidden_dropout
						snake_case__				:   Any     =						attention_dropout
						snake_case__				:   List[Any]     =						activation_dropout
						snake_case__				:   int     =						max_position_embeddings
						snake_case__				:   Any     =						type_vocab_size
						snake_case__				:   List[str]     =						2
						snake_case__				:   Union[str, Any]     =						num_labels
						snake_case__				:   Dict     =						num_choices
						snake_case__				:   Optional[int]     =						scope
						snake_case__				:   int     =						initializer_std

						# Used in the tests to check the size of the first attention layer
						snake_case__				:   Optional[Any]     =						n_head
						# Used in the tests to check the size of the first hidden state
						snake_case__				:   Union[str, Any]     =						self.d_model
						# Used in the tests to check the number of output hidden states/attentions
						snake_case__				:   Optional[Any]     =						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__				:   int     =						self.num_hidden_layers + 2



					def  _lowercase  (      self  :						List[str] ):
						snake_case__				:   int     =						ids_tensor([self.batch_size, self.seq_length]							,	self.vocab_size )

						snake_case__				:   Optional[int]     =						None
						if self.use_input_mask:
							snake_case__				:   Optional[Any]     =						random_attention_mask([self.batch_size, self.seq_length] )

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

						snake_case__				:   List[Any]     =						None
						snake_case__				:   str     =						None
						snake_case__				:   Tuple     =						None
						if self.use_labels:
							snake_case__				:   Tuple     =						ids_tensor([self.batch_size]							,	self.type_sequence_label_size )
							snake_case__				:   Dict     =						ids_tensor([self.batch_size, self.seq_length]							,	self.num_labels )
							snake_case__				:   Optional[Any]     =						ids_tensor([self.batch_size]							,	self.num_choices )

						snake_case__				:   Optional[int]     =						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							,	)

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



					def  _lowercase  (      self  :						Optional[Any]							,	__A  :						Optional[int]							,	__A  :						Any							,	__A  :						Tuple							,	__A  :						int							,	__A  :						Dict							,	__A  :						Union[str, Any]							,	__A  :						int							,	):
						snake_case__				:   int     =						TFFunnelModel(config=__A )
						snake_case__				:   List[Any]     =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						snake_case__				:   Tuple     =						model(__A )

						snake_case__				:   List[str]     =						[input_ids, input_mask]
						snake_case__				:   str     =						model(__A )

						snake_case__				:   Optional[int]     =						model(__A )
						self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.d_model) )

						snake_case__				:   Optional[int]     =						False
						snake_case__				:   Optional[int]     =						TFFunnelModel(config=__A )
						snake_case__				:   Dict     =						model(__A )
						self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.d_model) )

						snake_case__				:   Any     =						False
						snake_case__				:   Any     =						TFFunnelModel(config=__A )
						snake_case__				:   int     =						model(__A )
						self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.d_model) )



					def  _lowercase  (      self  :						Optional[int]							,	__A  :						Dict							,	__A  :						List[str]							,	__A  :						List[Any]							,	__A  :						Optional[int]							,	__A  :						Union[str, Any]							,	__A  :						int							,	__A  :						List[Any]							,	):
						snake_case__				:   Any     =						TFFunnelBaseModel(config=__A )
						snake_case__				:   List[Any]     =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						snake_case__				:   Optional[int]     =						model(__A )

						snake_case__				:   Optional[Any]     =						[input_ids, input_mask]
						snake_case__				:   List[str]     =						model(__A )

						snake_case__				:   Tuple     =						model(__A )
						self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, 2, self.d_model) )

						snake_case__				:   List[Any]     =						False
						snake_case__				:   Dict     =						TFFunnelBaseModel(config=__A )
						snake_case__				:   int     =						model(__A )
						self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, 3, self.d_model) )

						snake_case__				:   str     =						False
						snake_case__				:   Union[str, Any]     =						TFFunnelBaseModel(config=__A )
						snake_case__				:   int     =						model(__A )
						self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, 2, self.d_model) )



					def  _lowercase  (      self  :						List[Any]							,	__A  :						List[str]							,	__A  :						int							,	__A  :						Optional[int]							,	__A  :						Dict							,	__A  :						Any							,	__A  :						Optional[Any]							,	__A  :						Dict							,	):
						snake_case__				:   List[Any]     =						TFFunnelForPreTraining(config=__A )
						snake_case__				:   Tuple     =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						snake_case__				:   str     =						model(__A )
						self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length) )



					def  _lowercase  (      self  :						str							,	__A  :						Optional[Any]							,	__A  :						Optional[int]							,	__A  :						Optional[Any]							,	__A  :						Optional[int]							,	__A  :						Optional[int]							,	__A  :						Dict							,	__A  :						Optional[Any]							,	):
						snake_case__				:   Union[str, Any]     =						TFFunnelForMaskedLM(config=__A )
						snake_case__				:   str     =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						snake_case__				:   Any     =						model(__A )
						self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.vocab_size) )



					def  _lowercase  (      self  :						List[str]							,	__A  :						List[str]							,	__A  :						Union[str, Any]							,	__A  :						Optional[Any]							,	__A  :						List[Any]							,	__A  :						Union[str, Any]							,	__A  :						str							,	__A  :						Tuple							,	):
						snake_case__				:   Optional[Any]     =						self.num_labels
						snake_case__				:   str     =						TFFunnelForSequenceClassification(config=__A )
						snake_case__				:   Tuple     =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						snake_case__				:   Optional[Any]     =						model(__A )
						self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.num_labels) )



					def  _lowercase  (      self  :						Any							,	__A  :						Optional[Any]							,	__A  :						List[Any]							,	__A  :						Any							,	__A  :						Any							,	__A  :						Any							,	__A  :						Dict							,	__A  :						Optional[Any]							,	):
						snake_case__				:   Optional[int]     =						self.num_choices
						snake_case__				:   Union[str, Any]     =						TFFunnelForMultipleChoice(config=__A )
						snake_case__				:   List[str]     =						tf.tile(tf.expand_dims(__A							,	1 )							,	(1, self.num_choices, 1) )
						snake_case__				:   str     =						tf.tile(tf.expand_dims(__A							,	1 )							,	(1, self.num_choices, 1) )
						snake_case__				:   int     =						tf.tile(tf.expand_dims(__A							,	1 )							,	(1, self.num_choices, 1) )
						snake_case__				:   Optional[int]     =						{
						    "input_ids": multiple_choice_inputs_ids,
						    "attention_mask": multiple_choice_input_mask,
						    "token_type_ids": multiple_choice_token_type_ids,
						}
						snake_case__				:   Optional[int]     =						model(__A )
						self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.num_choices) )



					def  _lowercase  (      self  :						Optional[int]							,	__A  :						int							,	__A  :						Optional[int]							,	__A  :						Dict							,	__A  :						str							,	__A  :						List[Any]							,	__A  :						Optional[int]							,	__A  :						List[str]							,	):
						snake_case__				:   List[str]     =						self.num_labels
						snake_case__				:   str     =						TFFunnelForTokenClassification(config=__A )
						snake_case__				:   int     =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						snake_case__				:   Any     =						model(__A )
						self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.num_labels) )



					def  _lowercase  (      self  :						Optional[int]							,	__A  :						Tuple							,	__A  :						List[Any]							,	__A  :						Union[str, Any]							,	__A  :						Dict							,	__A  :						Dict							,	__A  :						Dict							,	__A  :						Dict							,	):
						snake_case__				:   int     =						TFFunnelForQuestionAnswering(config=__A )
						snake_case__				:   List[Any]     =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						snake_case__				:   int     =						model(__A )
						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  _lowercase  (      self  :						int ):
						snake_case__				:   Tuple     =						self.prepare_config_and_inputs()
						(
						    (
						    snake_case__
						),		(
						    snake_case__
						),		(
						    snake_case__
						),		(
						    snake_case__
						),		(
						    snake_case__
						),		(
						    snake_case__
						),		(
						    snake_case__
						),		
						)				:   Tuple     =						config_and_inputs
						snake_case__				:   Any     =						{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
						return config, inputs_dict



@require_tf
class 							SCREAMING_SNAKE_CASE__		(       UpperCamelCase_				,     UpperCamelCase_				,     unittest.TestCase		):





					"""simple docstring"""






					a_			=	(
					    (
					        TFFunnelModel,
					        TFFunnelForMaskedLM,
					        TFFunnelForPreTraining,
					        TFFunnelForQuestionAnswering,
					        TFFunnelForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					a_			=	(
					    {
					        "feature-extraction": (TFFunnelBaseModel, TFFunnelModel),
					        "fill-mask": TFFunnelForMaskedLM,
					        "question-answering": TFFunnelForQuestionAnswering,
					        "text-classification": TFFunnelForSequenceClassification,
					        "token-classification": TFFunnelForTokenClassification,
					        "zero-shot": TFFunnelForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					a_			=	False
					a_			=	False



					def  _lowercase  (      self  :						Dict ):
						snake_case__				:   Dict     =						TFFunnelModelTester(self )
						snake_case__				:   List[Any]     =						ConfigTester(self							,	config_class=__A )



					def  _lowercase  (      self  :						Union[str, Any] ):
						self.config_tester.run_common_tests()



					def  _lowercase  (      self  :						List[str] ):
						snake_case__				:   List[Any]     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__A )



					def  _lowercase  (      self  :						Optional[Any] ):
						snake_case__				:   Dict     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*__A )



					def  _lowercase  (      self  :						int ):
						snake_case__				:   int     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__A )



					def  _lowercase  (      self  :						Tuple ):
						snake_case__				:   Dict     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__A )






					def  _lowercase  (      self  :						Optional[Any] ):
						snake_case__				:   Optional[int]     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*__A )



@require_tf
class 							SCREAMING_SNAKE_CASE__		(       UpperCamelCase_				,     unittest.TestCase		):





					"""simple docstring"""






					a_			=	(
					    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
					)
					a_			=	False
					a_			=	False



					def  _lowercase  (      self  :						Any ):
						snake_case__				:   Any     =						TFFunnelModelTester(self							,	base=__A )
						snake_case__				:   Dict     =						ConfigTester(self							,	config_class=__A )



					def  _lowercase  (      self  :						List[str] ):
						self.config_tester.run_common_tests()



					def  _lowercase  (      self  :						Optional[Any] ):
						snake_case__				:   str     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_base_model(*__A )



					def  _lowercase  (      self  :						List[Any] ):
						snake_case__				:   Union[str, Any]     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*__A )






					def  _lowercase  (      self  :						Any ):
						snake_case__				:   List[str]     =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*__A )











# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
	import jax
	import jaxlib

__lowerCamelCase      :    Optional[int]   		=   get_logger()

__lowerCamelCase      :    Optional[dict]   		=   None



class 							SCREAMING_SNAKE_CASE__		(       TensorFormatter[Mapping, "jax.Array", Mapping]		):





					"""simple docstring"""






					def __init__(      self  :						Optional[Any]							,	__A  :						Dict=None							,	__A  :						List[str]=None							,	**__A  :						str ):
						super().__init__(features=__A )
						import jax
						from jaxlib.xla_client import Device

						if isinstance(__A							,	__A ):
							raise ValueError(
							    f'''Expected {device} to be a `str` not {type(__A )}, as `jaxlib.xla_extension.Device` '''
							    "is not serializable neither with `pickle` nor with `dill`. Instead you can surround "
							    "the device with `str()` to get its string identifier that will be internally mapped "
							    "to the actual `jaxlib.xla_extension.Device`." )
						snake_case__				:   List[Any]     =						device if isinstance(__A							,	__A ) else str(jax.devices()[0] )
						# using global variable since `jaxlib.xla_extension.Device` is not serializable neither
						# with `pickle` nor with `dill`, so we need to use a global variable instead
						global DEVICE_MAPPING
						if DEVICE_MAPPING is None:
							snake_case__				:   Any     =						self._map_devices_to_str()
						if self.device not in list(DEVICE_MAPPING.keys() ):
							logger.warning(
							    f'''Device with string identifier {self.device} not listed among the available '''
							    f'''devices: {list(DEVICE_MAPPING.keys() )}, so falling back to the default '''
							    f'''device: {str(jax.devices()[0] )}.''' )
							snake_case__				:   str     =						str(jax.devices()[0] )
						snake_case__				:   str     =						jnp_array_kwargs



					@staticmethod
					def  _lowercase  (      ):
						import jax

						return {str(__A ): device for device in jax.devices()}



					def  _lowercase  (      self  :						Optional[Any]							,	__A  :						str ):
						import jax
						import jax.numpy as jnp

						if isinstance(__A							,	__A ) and column:
							if all(
							    isinstance(__A							,	jax.Array ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column ):
								return jnp.stack(__A							,	axis=0 )
						return column



					def  _lowercase  (      self  :						int							,	__A  :						Tuple ):
						import jax
						import jax.numpy as jnp

						if isinstance(__A							,	(str, bytes, type(__A )) ):
							return value
						elif isinstance(__A							,	(np.character, np.ndarray) ) and np.issubdtype(value.dtype							,	np.character ):
							return value.tolist()

						snake_case__				:   Optional[int]     =						{}

						if isinstance(__A							,	(np.number, np.ndarray) ) and np.issubdtype(value.dtype							,	np.integer ):
							# the default int precision depends on the jax config
							# see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
							if jax.config.jax_enable_xaa:
								snake_case__				:   Any     =						{"dtype": jnp.intaa}
							else:
								snake_case__				:   Tuple     =						{"dtype": jnp.intaa}
						elif isinstance(__A							,	(np.number, np.ndarray) ) and np.issubdtype(value.dtype							,	np.floating ):
							snake_case__				:   str     =						{"dtype": jnp.floataa}
						elif config.PIL_AVAILABLE and "PIL" in sys.modules:
							import PIL.Image

							if isinstance(__A							,	PIL.Image.Image ):
								snake_case__				:   Optional[Any]     =						np.asarray(__A )

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
						global DEVICE_MAPPING
						if DEVICE_MAPPING is None:
							snake_case__				:   int     =						self._map_devices_to_str()

						with jax.default_device(DEVICE_MAPPING[self.device] ):
							# calling jnp.array on a np.ndarray does copy the data
							# see https://github.com/google/jax/issues/4486
							return jnp.array(__A							,	**{**default_dtype, **self.jnp_array_kwargs} )



					def  _lowercase  (      self  :						Union[str, Any]							,	__A  :						Optional[int] ):
						import jax

						# support for torch, tf, jax etc.
						if config.TORCH_AVAILABLE and "torch" in sys.modules:
							import torch

							if isinstance(__A							,	torch.Tensor ):
								return self._tensorize(data_struct.detach().cpu().numpy()[()] )
						if hasattr(__A							,	"__array__" ) and not isinstance(__A							,	jax.Array ):
							snake_case__				:   Union[str, Any]     =						data_struct.__array__()
						# support for nested types like struct of list of struct
						if isinstance(__A							,	np.ndarray ):
							if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
								return self._consolidate([self.recursive_tensorize(__A ) for substruct in data_struct] )
						elif isinstance(__A							,	(list, tuple) ):
							return self._consolidate([self.recursive_tensorize(__A ) for substruct in data_struct] )
						return self._tensorize(__A )



					def  _lowercase  (      self  :						Tuple							,	__A  :						dict ):
						return map_nested(self._recursive_tensorize							,	__A							,	map_list=__A )



					def  _lowercase  (      self  :						Optional[int]							,	__A  :						pa.Table ):
						snake_case__				:   int     =						self.numpy_arrow_extractor().extract_row(__A )
						snake_case__				:   Tuple     =						self.python_features_decoder.decode_row(__A )
						return self.recursive_tensorize(__A )



					def  _lowercase  (      self  :						Optional[Any]							,	__A  :						pa.Table ):
						snake_case__				:   Any     =						self.numpy_arrow_extractor().extract_column(__A )
						snake_case__				:   Optional[int]     =						self.python_features_decoder.decode_column(__A							,	pa_table.column_names[0] )
						snake_case__				:   List[Any]     =						self.recursive_tensorize(__A )
						snake_case__				:   Dict     =						self._consolidate(__A )
						return column






					def  _lowercase  (      self  :						str							,	__A  :						pa.Table ):
						snake_case__				:   Any     =						self.numpy_arrow_extractor().extract_batch(__A )
						snake_case__				:   int     =						self.python_features_decoder.decode_batch(__A )
						snake_case__				:   List[Any]     =						self.recursive_tensorize(__A )
						for column_name in batch:
							snake_case__				:   Any     =						self._consolidate(batch[column_name] )
						return batch



