import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_PRETRAINING_MAPPING,
		    NezhaForMaskedLM,
		    NezhaForMultipleChoice,
		    NezhaForNextSentencePrediction,
		    NezhaForPreTraining,
		    NezhaForQuestionAnswering,
		    NezhaForSequenceClassification,
		    NezhaForTokenClassification,
		    NezhaModel,
		)
		from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST




class    _UpperCamelCase      :


			"""simple docstring"""






			def __init__(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__=13      ,   lowerCAmelCase__=7      ,   lowerCAmelCase__=True      ,   lowerCAmelCase__=True      ,   lowerCAmelCase__=True      ,   lowerCAmelCase__=True      ,   lowerCAmelCase__=99      ,   lowerCAmelCase__=32      ,   lowerCAmelCase__=5      ,   lowerCAmelCase__=4      ,   lowerCAmelCase__=37      ,   lowerCAmelCase__="gelu"      ,   lowerCAmelCase__=0.1      ,   lowerCAmelCase__=0.1      ,   lowerCAmelCase__=1_28      ,   lowerCAmelCase__=32      ,   lowerCAmelCase__=16      ,   lowerCAmelCase__=2      ,   lowerCAmelCase__=0.02      ,   lowerCAmelCase__=3      ,   lowerCAmelCase__=4      ,   lowerCAmelCase__=None      ,   )     ->					str:

							'''simple docstring'''


							__lowercase 			=							parent
							__lowercase 			=							batch_size
							__lowercase 			=							seq_length
							__lowercase 			=							is_training
							__lowercase 			=							use_input_mask
							__lowercase 			=							use_token_type_ids
							__lowercase 			=							use_labels
							__lowercase 			=							vocab_size
							__lowercase 			=							hidden_size
							__lowercase 			=							num_hidden_layers
							__lowercase 			=							num_attention_heads
							__lowercase 			=							intermediate_size
							__lowercase 			=							hidden_act
							__lowercase 			=							hidden_dropout_prob
							__lowercase 			=							attention_probs_dropout_prob
							__lowercase 			=							max_position_embeddings
							__lowercase 			=							type_vocab_size
							__lowercase 			=							type_sequence_label_size
							__lowercase 			=							initializer_range
							__lowercase 			=							num_labels
							__lowercase 			=							num_choices
							__lowercase 			=							scope





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


							__lowercase 			=							ids_tensor([self.batch_size, self.seq_length]      ,   self.vocab_size						)

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

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

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

							__lowercase 			=							self.get_config()

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





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

							'''simple docstring'''


							return NezhaConfig(
							    vocab_size=self.vocab_size      ,   hidden_size=self.hidden_size      ,   num_hidden_layers=self.num_hidden_layers      ,   num_attention_heads=self.num_attention_heads      ,   intermediate_size=self.intermediate_size      ,   hidden_act=self.hidden_act      ,   hidden_dropout_prob=self.hidden_dropout_prob      ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,   max_position_embeddings=self.max_position_embeddings      ,   type_vocab_size=self.type_vocab_size      ,   is_decoder=lowerCAmelCase__      ,   initializer_range=self.initializer_range      ,   )





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

							'''simple docstring'''


							(
							    (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     
							) 			=							self.prepare_config_and_inputs()

							__lowercase 			=							True
							__lowercase 			=							floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
							__lowercase 			=							ids_tensor([self.batch_size, self.seq_length]      ,   vocab_size=2						)

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





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)     ->					List[Any]:

							'''simple docstring'''


							__lowercase 			=							NezhaModel(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__						)
							__lowercase 			=							model(lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__						)
							__lowercase 			=							model(lowerCAmelCase__						)
							self.parent.assertEqual(result.last_hidden_state.shape      ,   (self.batch_size, self.seq_length, self.hidden_size)						)
							self.parent.assertEqual(result.pooler_output.shape      ,   (self.batch_size, self.hidden_size)						)





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   )     ->					Any:

							'''simple docstring'''


							__lowercase 			=							True
							__lowercase 			=							NezhaModel(lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(
							    lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   encoder_hidden_states=lowerCAmelCase__      ,   encoder_attention_mask=lowerCAmelCase__      ,   )
							__lowercase 			=							model(
							    lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   encoder_hidden_states=lowerCAmelCase__      ,   )
							__lowercase 			=							model(lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__						)
							self.parent.assertEqual(result.last_hidden_state.shape      ,   (self.batch_size, self.seq_length, self.hidden_size)						)
							self.parent.assertEqual(result.pooler_output.shape      ,   (self.batch_size, self.hidden_size)						)





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)     ->					Optional[Any]:

							'''simple docstring'''


							__lowercase 			=							NezhaForMaskedLM(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   labels=lowerCAmelCase__						)
							self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size)						)





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)     ->					str:

							'''simple docstring'''


							__lowercase 			=							NezhaForNextSentencePrediction(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(
							    lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   labels=lowerCAmelCase__      ,   )
							self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, 2)						)





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)     ->					Tuple:

							'''simple docstring'''


							__lowercase 			=							NezhaForPreTraining(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(
							    lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   labels=lowerCAmelCase__      ,   next_sentence_label=lowerCAmelCase__      ,   )
							self.parent.assertEqual(result.prediction_logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size)						)
							self.parent.assertEqual(result.seq_relationship_logits.shape      ,   (self.batch_size, 2)						)





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)     ->					Optional[Any]:

							'''simple docstring'''


							__lowercase 			=							NezhaForQuestionAnswering(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(
							    lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   start_positions=lowerCAmelCase__      ,   end_positions=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        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)     ->					List[str]:

							'''simple docstring'''


							__lowercase 			=							self.num_labels
							__lowercase 			=							NezhaForSequenceClassification(lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   labels=lowerCAmelCase__						)
							self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.num_labels)						)





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

							'''simple docstring'''


							__lowercase 			=							self.num_labels
							__lowercase 			=							NezhaForTokenClassification(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							model(lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   labels=lowerCAmelCase__						)
							self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.num_labels)						)





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)     ->					Any:

							'''simple docstring'''


							__lowercase 			=							self.num_choices
							__lowercase 			=							NezhaForMultipleChoice(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase 			=							input_ids.unsqueeze(1						).expand(-1      ,   self.num_choices      ,   -1						).contiguous()
							__lowercase 			=							token_type_ids.unsqueeze(1						).expand(-1      ,   self.num_choices      ,   -1						).contiguous()
							__lowercase 			=							input_mask.unsqueeze(1						).expand(-1      ,   self.num_choices      ,   -1						).contiguous()
							__lowercase 			=							model(
							    lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__      ,   token_type_ids=lowerCAmelCase__      ,   labels=lowerCAmelCase__      ,   )
							self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.num_choices)						)





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

							'''simple docstring'''


							__lowercase 			=							self.prepare_config_and_inputs()
							(
							    (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     
							) 			=							config_and_inputs
							__lowercase 			=							{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
							return config, inputs_dict




@require_torch
class    _UpperCamelCase      (						_UpperCAmelCase						,_UpperCAmelCase						,_UpperCAmelCase						,unittest.TestCase ):


			"""simple docstring"""






			__a  : str			 =						(
			    (
			        NezhaModel,
			        NezhaForMaskedLM,
			        NezhaForMultipleChoice,
			        NezhaForNextSentencePrediction,
			        NezhaForPreTraining,
			        NezhaForQuestionAnswering,
			        NezhaForSequenceClassification,
			        NezhaForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			__a  : List[Any]			 =						(
			    {
			        '''feature-extraction''': NezhaModel,
			        '''fill-mask''': NezhaForMaskedLM,
			        '''question-answering''': NezhaForQuestionAnswering,
			        '''text-classification''': NezhaForSequenceClassification,
			        '''token-classification''': NezhaForTokenClassification,
			        '''zero-shot''': NezhaForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			__a  : Dict			 =						True





			def        _SCREAMING_SNAKE_CASE			(   self      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__=False						)     ->					List[str]:

							'''simple docstring'''


							__lowercase 			=							super()._prepare_for_class(lowerCAmelCase__      ,   lowerCAmelCase__      ,   return_labels=lowerCAmelCase__						)

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





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

							'''simple docstring'''


							__lowercase 			=							NezhaModelTester(self						)
							__lowercase 			=							ConfigTester(self      ,   config_class=lowerCAmelCase__      ,   hidden_size=37						)





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

							'''simple docstring'''


							self.config_tester.run_common_tests()





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					int:

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCAmelCase__						)





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

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase__						)





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

							'''simple docstring'''


							(
							    (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     (
							    __lowercase
							)			,     
							) 			=							self.model_tester.prepare_config_and_inputs_for_decoder()

							__lowercase 			=							None

							self.model_tester.create_and_check_model_as_decoder(
							    lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__      ,   )





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

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Tuple:

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Dict:

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_next_sequence_prediction(*lowerCAmelCase__						)





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

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Dict:

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase__						)





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

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


							__lowercase 			=							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__						)





			@slow
			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


							for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__lowercase 			=							NezhaModel.from_pretrained(lowerCAmelCase__						)
											self.assertIsNotNone(lowerCAmelCase__						)





			@slow
			@require_torch_gpu
			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Any:

							'''simple docstring'''


							__lowercase			,     __lowercase 			=							self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes:
											# NezhaForMultipleChoice behaves incorrectly in JIT environments.
											if model_class == NezhaForMultipleChoice:
															return

											__lowercase 			=							True
											__lowercase 			=							model_class(config=lowerCAmelCase__						)

											__lowercase 			=							self._prepare_for_class(lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							torch.jit.trace(
											    lowerCAmelCase__      ,   (inputs_dict['''input_ids'''].to('''cpu'''						), inputs_dict['''attention_mask'''].to('''cpu'''						))						)

											with tempfile.TemporaryDirectory() as tmp:
															torch.jit.save(lowerCAmelCase__      ,   os.path.join(lowerCAmelCase__      ,   '''bert.pt'''						)						)
															__lowercase 			=							torch.jit.load(os.path.join(lowerCAmelCase__      ,   '''bert.pt'''						)      ,   map_location=lowerCAmelCase__						)
															loaded(inputs_dict['''input_ids'''].to(lowerCAmelCase__						)      ,   inputs_dict['''attention_mask'''].to(lowerCAmelCase__						)						)




@require_torch
class    _UpperCamelCase      (						unittest.TestCase ):


			"""simple docstring"""






			@slow
			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


							__lowercase 			=							NezhaModel.from_pretrained('''sijunhe/nezha-cn-base'''						)
							__lowercase 			=							torch.tensor([[0, 1, 2, 3, 4, 5]]						)
							__lowercase 			=							torch.tensor([[0, 1, 1, 1, 1, 1]]						)
							with torch.no_grad():
											__lowercase 			=							model(lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__						)[0]
							__lowercase 			=							torch.Size((1, 6, 7_68)						)
							self.assertEqual(output.shape      ,   lowerCAmelCase__						)
							__lowercase 			=							torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]						)

							self.assertTrue(torch.allclose(output[:, 1:4, 1:4]      ,   lowerCAmelCase__      ,   atol=1E-4						)						)





			@slow
			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					List[str]:

							'''simple docstring'''


							__lowercase 			=							NezhaForMaskedLM.from_pretrained('''sijunhe/nezha-cn-base'''						)
							__lowercase 			=							torch.tensor([[0, 1, 2, 3, 4, 5]]						)
							__lowercase 			=							torch.tensor([[1, 1, 1, 1, 1, 1]]						)
							with torch.no_grad():
											__lowercase 			=							model(lowerCAmelCase__      ,   attention_mask=lowerCAmelCase__						)[0]
							__lowercase 			=							torch.Size((1, 6, 2_11_28)						)
							self.assertEqual(output.shape      ,   lowerCAmelCase__						)
							__lowercase 			=							torch.tensor(
							    [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]						)

							self.assertTrue(torch.allclose(output[:, 1:4, 1:4]      ,   lowerCAmelCase__      ,   atol=1E-4						)						)
from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


__a  :     Optional[Any]       =  logging.get_logger(__name__)


__a  :     List[str]       =  TypeVar("""DatasetType""", Dataset, IterableDataset)





def 				UpperCAmelCase						(    lowercase							,			lowercase = None							,			lowercase = None							,			lowercase = None							,			lowercase = None							,			lowercase = "first_exhausted"							,			):







				"""simple docstring"""



				from .arrow_dataset import Dataset
				from .iterable_dataset import IterableDataset

				if not datasets:
								raise ValueError('''Unable to interleave an empty list of datasets.'''							)
				for i, dataset in enumerate(lowercase							):
								if not isinstance(lowercase							,			(Dataset, IterableDataset)							):
												if isinstance(lowercase							,			(DatasetDict, IterableDatasetDict)							):
																if not dataset:
																				raise ValueError(
																				    F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
																				    '''is an empty dataset dictionary.'''							)
																raise ValueError(
																    F"Dataset at position {i} has at least one split: {list(lowercase							)}\n"
																    F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(lowercase							)							)}']"							)
												raise ValueError(
												    F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(lowercase							).__name__}."							)
								if i == 0:
												__lowercase			,     __lowercase 			=							(
												    (Dataset, IterableDataset) if isinstance(lowercase							,			lowercase							) else (IterableDataset, Dataset)
												)
								elif not isinstance(lowercase							,			lowercase							):
												raise ValueError(
												    F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."							)
				if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
								raise ValueError(F"{stopping_strategy} is not supported. Please enter a valid stopping_strategy."							)
				if dataset_type is Dataset:
								return _interleave_map_style_datasets(
								    lowercase							,			lowercase							,			lowercase							,			info=lowercase							,			split=lowercase							,			stopping_strategy=lowercase							)
				else:
								return _interleave_iterable_datasets(
								    lowercase							,			lowercase							,			lowercase							,			info=lowercase							,			split=lowercase							,			stopping_strategy=lowercase							)





def 				UpperCAmelCase						(    lowercase							,			lowercase = None							,			lowercase = None							,			lowercase = 0							,			):







				"""simple docstring"""



				if not dsets:
								raise ValueError('''Unable to concatenate an empty list of datasets.'''							)
				for i, dataset in enumerate(lowercase							):
								if not isinstance(lowercase							,			(Dataset, IterableDataset)							):
												if isinstance(lowercase							,			(DatasetDict, IterableDatasetDict)							):
																if not dataset:
																				raise ValueError(
																				    F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
																				    '''is an empty dataset dictionary.'''							)
																raise ValueError(
																    F"Dataset at position {i} has at least one split: {list(lowercase							)}\n"
																    F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(lowercase							)							)}']"							)
												raise ValueError(
												    F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(lowercase							).__name__}."							)
								if i == 0:
												__lowercase			,     __lowercase 			=							(
												    (Dataset, IterableDataset) if isinstance(lowercase							,			lowercase							) else (IterableDataset, Dataset)
												)
								elif not isinstance(lowercase							,			lowercase							):
												raise ValueError(
												    F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."							)
				if dataset_type is Dataset:
								return _concatenate_map_style_datasets(lowercase							,			info=lowercase							,			split=lowercase							,			axis=lowercase							)
				else:
								return _concatenate_iterable_datasets(lowercase							,			info=lowercase							,			split=lowercase							,			axis=lowercase							)