



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



	def __init__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__=13			,	UpperCAmelCase__=7			,	UpperCAmelCase__=True			,	UpperCAmelCase__=True			,	UpperCAmelCase__=True			,	UpperCAmelCase__=True			,	UpperCAmelCase__=99			,	UpperCAmelCase__=[1, 1, 2]			,	UpperCAmelCase__=1			,	UpperCAmelCase__=32			,	UpperCAmelCase__=4			,	UpperCAmelCase__=8			,	UpperCAmelCase__=37			,	UpperCAmelCase__="gelu_new"			,	UpperCAmelCase__=0.1			,	UpperCAmelCase__=0.1			,	UpperCAmelCase__=0.0			,	UpperCAmelCase__=512			,	UpperCAmelCase__=3			,	UpperCAmelCase__=0.02			,	UpperCAmelCase__=3			,	UpperCAmelCase__=4			,	UpperCAmelCase__=None			,	UpperCAmelCase__=False			,	):
							A__							=       parent
							A__							=       batch_size
							A__							=       seq_length
							A__							=       is_training
							A__							=       use_input_mask
							A__							=       use_token_type_ids
							A__							=       use_labels
							A__							=       vocab_size
							A__							=       block_sizes
							A__							=       num_decoder_layers
							A__							=       d_model
							A__							=       n_head
							A__							=       d_head
							A__							=       d_inner
							A__							=       hidden_act
							A__							=       hidden_dropout
							A__							=       attention_dropout
							A__							=       activation_dropout
							A__							=       max_position_embeddings
							A__							=       type_vocab_size
							A__							=       2
							A__							=       num_labels
							A__							=       num_choices
							A__							=       scope
							A__							=       initializer_std

							# Used in the tests to check the size of the first attention layer
							A__							=       n_head
							# Used in the tests to check the size of the first hidden state
							A__							=       self.d_model
							# Used in the tests to check the number of output hidden states/attentions
							A__							=       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:
													A__							=       self.num_hidden_layers + 2



	def 			__A							(       self		):
							A__							=       ids_tensor([self.batch_size, self.seq_length]			,	self.vocab_size		)

							A__							=       None
							if self.use_input_mask:
													A__							=       random_attention_mask([self.batch_size, self.seq_length]		)

							A__							=       None
							if self.use_token_type_ids:
													A__							=       ids_tensor([self.batch_size, self.seq_length]			,	self.type_vocab_size		)

							A__							=       None
							A__							=       None
							A__							=       None
							if self.use_labels:
													A__							=       ids_tensor([self.batch_size]			,	self.type_sequence_label_size		)
													A__							=       ids_tensor([self.batch_size, self.seq_length]			,	self.num_labels		)
													A__							=       ids_tensor([self.batch_size]			,	self.num_choices		)

							A__							=       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 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       TFFunnelModel(config=__lowerCAmelCase		)
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							A__							=       model(__lowerCAmelCase		)

							A__							=       [input_ids, input_mask]
							A__							=       model(__lowerCAmelCase		)

							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, self.seq_length, self.d_model)		)

							A__							=       False
							A__							=       TFFunnelModel(config=__lowerCAmelCase		)
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, self.seq_length, self.d_model)		)

							A__							=       False
							A__							=       TFFunnelModel(config=__lowerCAmelCase		)
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, self.seq_length, self.d_model)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       TFFunnelBaseModel(config=__lowerCAmelCase		)
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							A__							=       model(__lowerCAmelCase		)

							A__							=       [input_ids, input_mask]
							A__							=       model(__lowerCAmelCase		)

							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, 2, self.d_model)		)

							A__							=       False
							A__							=       TFFunnelBaseModel(config=__lowerCAmelCase		)
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, 3, self.d_model)		)

							A__							=       False
							A__							=       TFFunnelBaseModel(config=__lowerCAmelCase		)
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, 2, self.d_model)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       TFFunnelForPreTraining(config=__lowerCAmelCase		)
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       TFFunnelForMaskedLM(config=__lowerCAmelCase		)
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.vocab_size)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       self.num_labels
							A__							=       TFFunnelForSequenceClassification(config=__lowerCAmelCase		)
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.num_labels)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       self.num_choices
							A__							=       TFFunnelForMultipleChoice(config=__lowerCAmelCase		)
							A__							=       tf.tile(tf.expand_dims(__lowerCAmelCase			,	1		)			,	(1, self.num_choices, 1)		)
							A__							=       tf.tile(tf.expand_dims(__lowerCAmelCase			,	1		)			,	(1, self.num_choices, 1)		)
							A__							=       tf.tile(tf.expand_dims(__lowerCAmelCase			,	1		)			,	(1, self.num_choices, 1)		)
							A__							=       {
							    "input_ids": multiple_choice_inputs_ids,
							    "attention_mask": multiple_choice_input_mask,
							    "token_type_ids": multiple_choice_token_type_ids,
							}
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.num_choices)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       self.num_labels
							A__							=       TFFunnelForTokenClassification(config=__lowerCAmelCase		)
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							A__							=       model(__lowerCAmelCase		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.num_labels)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       TFFunnelForQuestionAnswering(config=__lowerCAmelCase		)
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							A__							=       model(__lowerCAmelCase		)
							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 			__A							(       self		):
							A__							=       self.prepare_config_and_inputs()
							(
							    (
							    A__       
							)						,       (
							    A__       
							)						,       (
							    A__       
							)						,       (
							    A__       
							)						,       (
							    A__       
							)						,       (
							    A__       
							)						,       (
							    A__       
							)						,              
							)							=       config_and_inputs
							A__							=       {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
							return config, inputs_dict






@require_tf
class      UpperCamelCase    (						_a       ,    _a       ,    unittest.TestCase     ):
	lowerCAmelCase	:   str				=	(
	    (
	        TFFunnelModel,
	        TFFunnelForMaskedLM,
	        TFFunnelForPreTraining,
	        TFFunnelForQuestionAnswering,
	        TFFunnelForTokenClassification,
	    )
	    if is_tf_available()
	    else ()
	)
	lowerCAmelCase	:   List[Any]				=	(
	    {
	        """feature-extraction""": (TFFunnelBaseModel, TFFunnelModel),
	        """fill-mask""": TFFunnelForMaskedLM,
	        """question-answering""": TFFunnelForQuestionAnswering,
	        """text-classification""": TFFunnelForSequenceClassification,
	        """token-classification""": TFFunnelForTokenClassification,
	        """zero-shot""": TFFunnelForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	lowerCAmelCase	:   List[Any]				=	False
	lowerCAmelCase	:   List[Any]				=	False



	def 			__A							(       self		):
							A__							=       TFFunnelModelTester(self		)
							A__							=       ConfigTester(self			,	config_class=__lowerCAmelCase		)



	def 			__A							(       self		):
							self.config_tester.run_common_tests()



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__lowerCAmelCase		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*__lowerCAmelCase		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase		)






@require_tf
class      UpperCamelCase    (						_a       ,    unittest.TestCase     ):
	lowerCAmelCase	:   Tuple				=	(
	    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
	)
	lowerCAmelCase	:   int				=	False
	lowerCAmelCase	:   Optional[Any]				=	False



	def 			__A							(       self		):
							A__							=       TFFunnelModelTester(self			,	base=__lowerCAmelCase		)
							A__							=       ConfigTester(self			,	config_class=__lowerCAmelCase		)



	def 			__A							(       self		):
							self.config_tester.run_common_tests()



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_base_model(*__lowerCAmelCase		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__lowerCAmelCase		)










import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def __init__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = False			,	UpperCAmelCase__ = False			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	**UpperCAmelCase__			,	):
							super().__init__(
							    UpperCAmelCase__			,	split=UpperCAmelCase__			,	features=UpperCAmelCase__			,	cache_dir=UpperCAmelCase__			,	keep_in_memory=UpperCAmelCase__			,	streaming=UpperCAmelCase__			,	num_proc=UpperCAmelCase__			,	**UpperCAmelCase__			,	)
							A__							=       field
							A__							=       path_or_paths if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		) else {self.split: path_or_paths}
							A__							=       Json(
							    cache_dir=UpperCAmelCase__			,	data_files=UpperCAmelCase__			,	features=UpperCAmelCase__			,	field=UpperCAmelCase__			,	**UpperCAmelCase__			,	)



	def 			__A							(       self		):
							# Build iterable dataset
							if self.streaming:
													A__							=       self.builder.as_streaming_dataset(split=self.split		)
							# Build regular (map-style) dataset
							else:
													A__							=       None
													A__							=       None
													A__							=       None
													A__							=       None

													self.builder.download_and_prepare(
													    download_config=UpperCAmelCase__			,	download_mode=UpperCAmelCase__			,	verification_mode=UpperCAmelCase__			,	base_path=UpperCAmelCase__			,	num_proc=self.num_proc			,	)
													A__							=       self.builder.as_dataset(
													    split=self.split			,	verification_mode=UpperCAmelCase__			,	in_memory=self.keep_in_memory		)
							return dataset






class      UpperCamelCase    :



	def __init__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	**UpperCAmelCase__			,	):
							if num_proc is not None and num_proc <= 0:
													raise ValueError(F"""num_proc {num_proc} must be an integer > 0."""		)

							A__							=       dataset
							A__							=       path_or_buf
							A__							=       batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
							A__							=       num_proc
							A__							=       "utf-8"
							A__							=       to_json_kwargs



	def 			__A							(       self		):
							A__							=       self.to_json_kwargs.pop("path_or_buf"			,	UpperCAmelCase__		)
							A__							=       self.to_json_kwargs.pop("orient"			,	"records"		)
							A__							=       self.to_json_kwargs.pop("lines"			,	True if orient == "records" else False		)
							A__							=       self.to_json_kwargs.pop("index"			,	False if orient in ["split", "table"] else True		)
							A__							=       self.to_json_kwargs.pop("compression"			,	UpperCAmelCase__		)

							if compression not in [None, "infer", "gzip", "bz2", "xz"]:
													raise NotImplementedError(F"""`datasets` currently does not support {compression} compression"""		)

							if isinstance(self.path_or_buf			,	(str, bytes, os.PathLike)		):
													with fsspec.open(self.path_or_buf			,	"wb"			,	compression=UpperCAmelCase__		) as buffer:
																			A__							=       self._write(file_obj=UpperCAmelCase__			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**self.to_json_kwargs		)
							else:
													if compression:
																			raise NotImplementedError(
																			    F"""The compression parameter is not supported when writing to a buffer, but compression={compression}"""
																			    " was passed. Please provide a local path instead."		)
													A__							=       self._write(
													    file_obj=self.path_or_buf			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**self.to_json_kwargs		)
							return written



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__						,       A__						,       A__						,       A__						,       A__							=       args

							A__							=       query_table(
							    table=self.dataset.data			,	key=slice(UpperCAmelCase__			,	offset + self.batch_size		)			,	indices=self.dataset._indices			,	)
							A__							=       batch.to_pandas().to_json(
							    path_or_buf=UpperCAmelCase__			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**UpperCAmelCase__		)
							if not json_str.endswith("\n"		):
													json_str += "\n"
							return json_str.encode(self.encoding		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__			,	):
							A__							=       0

							if self.num_proc is None or self.num_proc == 1:
													for offset in logging.tqdm(
													    range(0			,	len(self.dataset		)			,	self.batch_size		)			,	unit="ba"			,	disable=not logging.is_progress_bar_enabled()			,	desc="Creating json from Arrow format"			,	):
																			A__							=       self._batch_json((offset, orient, lines, index, to_json_kwargs)		)
																			written += file_obj.write(UpperCAmelCase__		)
							else:
													A__						,       A__							=       len(self.dataset		), self.batch_size
													with multiprocessing.Pool(self.num_proc		) as pool:
																			for json_str in logging.tqdm(
																			    pool.imap(
																			        self._batch_json			,	[(offset, orient, lines, index, to_json_kwargs) for offset in range(0			,	UpperCAmelCase__			,	UpperCAmelCase__		)]			,	)			,	total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size			,	unit="ba"			,	disable=not logging.is_progress_bar_enabled()			,	desc="Creating json from Arrow format"			,	):
																									written += file_obj.write(UpperCAmelCase__		)

							return written





