






"""simple docstring"""



from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow

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 numpy as np
						import tensorflow as tf

						from transformers import (
						    TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    FlaubertConfig,
						    TFFlaubertForMultipleChoice,
						    TFFlaubertForQuestionAnsweringSimple,
						    TFFlaubertForSequenceClassification,
						    TFFlaubertForTokenClassification,
						    TFFlaubertModel,
						    TFFlaubertWithLMHeadModel,
						)





class 					_UpperCAmelCase       :




					def __init__(				self     :						List[str]    ,						A     :						Any    ,						)  ->  Dict:
							lowercase_ : Tuple							       =   parent
							lowercase_ : str							       =   13
							lowercase_ : Optional[Any]							       =   7
							lowercase_ : Any							       =   True
							lowercase_ : str							       =   True
							lowercase_ : List[str]							       =   True
							lowercase_ : int							       =   True
							lowercase_ : Dict							       =   True
							lowercase_ : int							       =   False
							lowercase_ : Dict							       =   False
							lowercase_ : Union[str, Any]							       =   False
							lowercase_ : List[Any]							       =   2
							lowercase_ : Optional[int]							       =   99
							lowercase_ : List[Any]							       =   0
							lowercase_ : Dict							       =   32
							lowercase_ : List[Any]							       =   2
							lowercase_ : Tuple							       =   4
							lowercase_ : List[Any]							       =   0.1
							lowercase_ : List[Any]							       =   0.1
							lowercase_ : Optional[Any]							       =   5_12
							lowercase_ : Optional[Any]							       =   16
							lowercase_ : List[str]							       =   2
							lowercase_ : str							       =   0.02
							lowercase_ : Any							       =   3
							lowercase_ : List[str]							       =   4
							lowercase_ : Dict							       =   '''last'''
							lowercase_ : int							       =   True
							lowercase_ : str							       =   None
							lowercase_ : Dict							       =   0




					def   A  (				self     :						List[str]   )  ->  List[Any]:
							lowercase_ : Union[str, Any]							       =   ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size   )
							lowercase_ : List[Any]							       =   random_attention_mask([self.batch_size, self.seq_length]    ,						dtype=tf.floataa   )

							lowercase_ : Any							       =   None
							if self.use_input_lengths:
									lowercase_ : Optional[Any]							       =   (
									    ids_tensor([self.batch_size]    ,						vocab_size=2   ) + self.seq_length - 2
									)  # small variation of seq_length

							lowercase_ : List[Any]							       =   None
							if self.use_token_type_ids:
									lowercase_ : Optional[Any]							       =   ids_tensor([self.batch_size, self.seq_length]    ,						self.n_langs   )

							lowercase_ : Optional[int]							       =   None
							lowercase_ : str							       =   None
							lowercase_ : Optional[Any]							       =   None
							if self.use_labels:
									lowercase_ : str							       =   ids_tensor([self.batch_size]    ,						self.type_sequence_label_size   )
									lowercase_ : Tuple							       =   ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels   )
									lowercase_ : Optional[int]							       =   ids_tensor([self.batch_size]    ,						2    ,						dtype=tf.floataa   )
									lowercase_ : List[Any]							       =   ids_tensor([self.batch_size]    ,						self.num_choices   )

							lowercase_ : Any							       =   FlaubertConfig(
							    vocab_size=self.vocab_size    ,						n_special=self.n_special    ,						emb_dim=self.hidden_size    ,						n_layers=self.num_hidden_layers    ,						n_heads=self.num_attention_heads    ,						dropout=self.hidden_dropout_prob    ,						attention_dropout=self.attention_probs_dropout_prob    ,						gelu_activation=self.gelu_activation    ,						sinusoidal_embeddings=self.sinusoidal_embeddings    ,						asm=self.asm    ,						causal=self.causal    ,						n_langs=self.n_langs    ,						max_position_embeddings=self.max_position_embeddings    ,						initializer_range=self.initializer_range    ,						summary_type=self.summary_type    ,						use_proj=self.use_proj    ,						bos_token_id=self.bos_token_id    ,						)

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




					def   A  (				self     :						int    ,						A     :						Dict    ,						A     :						Dict    ,						A     :						str    ,						A     :						Tuple    ,						A     :						Optional[Any]    ,						A     :						str    ,						A     :						Union[str, Any]    ,						A     :						Dict    ,						A     :						Optional[int]    ,						)  ->  int:
							lowercase_ : str							       =   TFFlaubertModel(config=A   )
							lowercase_ : int							       =   {'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
							lowercase_ : int							       =   model(A   )

							lowercase_ : Any							       =   [input_ids, input_mask]
							lowercase_ : Optional[int]							       =   model(A   )
							self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)   )




					def   A  (				self     :						Optional[Any]    ,						A     :						int    ,						A     :						List[str]    ,						A     :						Dict    ,						A     :						int    ,						A     :						int    ,						A     :						Dict    ,						A     :						Optional[int]    ,						A     :						Dict    ,						A     :						int    ,						)  ->  Optional[Any]:
							lowercase_ : int							       =   TFFlaubertWithLMHeadModel(A   )

							lowercase_ : List[Any]							       =   {'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
							lowercase_ : List[str]							       =   model(A   )

							self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)   )




					def   A  (				self     :						List[str]    ,						A     :						Optional[int]    ,						A     :						Tuple    ,						A     :						Dict    ,						A     :						List[Any]    ,						A     :						Dict    ,						A     :						Any    ,						A     :						List[Any]    ,						A     :						List[Any]    ,						A     :						str    ,						)  ->  Union[str, Any]:
							lowercase_ : List[Any]							       =   TFFlaubertForQuestionAnsweringSimple(A   )

							lowercase_ : Tuple							       =   {'''input_ids''': input_ids, '''lengths''': input_lengths}

							lowercase_ : List[str]							       =   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   A  (				self     :						int    ,						A     :						Optional[Any]    ,						A     :						int    ,						A     :						List[str]    ,						A     :						Optional[Any]    ,						A     :						Tuple    ,						A     :						Dict    ,						A     :						Any    ,						A     :						Any    ,						A     :						str    ,						)  ->  Optional[int]:
							lowercase_ : str							       =   TFFlaubertForSequenceClassification(A   )

							lowercase_ : Any							       =   {'''input_ids''': input_ids, '''lengths''': input_lengths}

							lowercase_ : Union[str, Any]							       =   model(A   )

							self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.type_sequence_label_size)   )




					def   A  (				self     :						List[Any]    ,						A     :						Tuple    ,						A     :						int    ,						A     :						Dict    ,						A     :						Any    ,						A     :						int    ,						A     :						Optional[int]    ,						A     :						str    ,						A     :						str    ,						A     :						int    ,						)  ->  Optional[int]:
							lowercase_ : Optional[Any]							       =   self.num_labels
							lowercase_ : List[Any]							       =   TFFlaubertForTokenClassification(config=A   )
							lowercase_ : int							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
							lowercase_ : int							       =   model(A   )
							self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)   )




					def   A  (				self     :						List[str]    ,						A     :						str    ,						A     :						Any    ,						A     :						int    ,						A     :						Dict    ,						A     :						Tuple    ,						A     :						List[Any]    ,						A     :						Optional[Any]    ,						A     :						List[Any]    ,						A     :						Optional[int]    ,						)  ->  Union[str, Any]:
							lowercase_ : Union[str, Any]							       =   self.num_choices
							lowercase_ : Union[str, Any]							       =   TFFlaubertForMultipleChoice(config=A   )
							lowercase_ : Union[str, Any]							       =   tf.tile(tf.expand_dims(A    ,						1   )    ,						(1, self.num_choices, 1)   )
							lowercase_ : Any							       =   tf.tile(tf.expand_dims(A    ,						1   )    ,						(1, self.num_choices, 1)   )
							lowercase_ : Dict							       =   tf.tile(tf.expand_dims(A    ,						1   )    ,						(1, self.num_choices, 1)   )
							lowercase_ : Any							       =   {
							    '''input_ids''': multiple_choice_inputs_ids,
							    '''attention_mask''': multiple_choice_input_mask,
							    '''token_type_ids''': multiple_choice_token_type_ids,
							}
							lowercase_ : Any							       =   model(A   )
							self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)   )






					def   A  (				self     :						Optional[int]   )  ->  str:
							lowercase_ : List[Any]							       =   self.prepare_config_and_inputs()
							(
							    (
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					(
							    lowercase_
							)       ,					
							) : Dict							       =   config_and_inputs
							lowercase_ : Optional[int]							       =   {
							    '''input_ids''': input_ids,
							    '''token_type_ids''': token_type_ids,
							    '''langs''': token_type_ids,
							    '''lengths''': input_lengths,
							}
							return config, inputs_dict





@require_tf
class 					_UpperCAmelCase       ( _A   ,						_A   ,						unittest.TestCase							):
					SCREAMING_SNAKE_CASE_		:				Union[str, Any]							    =   (
					    (
					        TFFlaubertModel,
					        TFFlaubertWithLMHeadModel,
					        TFFlaubertForSequenceClassification,
					        TFFlaubertForQuestionAnsweringSimple,
					        TFFlaubertForTokenClassification,
					        TFFlaubertForMultipleChoice,
					    )
					    if is_tf_available()
					    else ()
					)
					SCREAMING_SNAKE_CASE_		:				Dict							    =   (
					    (TFFlaubertWithLMHeadModel,) if is_tf_available() else ()
					)  # TODO (PVP): Check other models whether language generation is also applicable
					SCREAMING_SNAKE_CASE_		:				Union[str, Any]							    =   (
					    {
					        "feature-extraction": TFFlaubertModel,
					        "fill-mask": TFFlaubertWithLMHeadModel,
					        "question-answering": TFFlaubertForQuestionAnsweringSimple,
					        "text-classification": TFFlaubertForSequenceClassification,
					        "token-classification": TFFlaubertForTokenClassification,
					        "zero-shot": TFFlaubertForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					SCREAMING_SNAKE_CASE_		:				str							    =   False
					SCREAMING_SNAKE_CASE_		:				Optional[Any]							    =   False




					def   A  (				self     :						Any    ,						A     :						Any    ,						A     :						Union[str, Any]    ,						A     :						Optional[int]    ,						A     :						int    ,						A     :						str   )  ->  Any:
							if (
							    pipeline_test_casse_name == "QAPipelineTests"
							    and tokenizer_name is not None
							    and not tokenizer_name.endswith('''Fast'''   )
							):
									# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
									# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
									# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
									return True

							return False




					def   A  (				self     :						Any   )  ->  Optional[int]:
							lowercase_ : Dict							       =   TFFlaubertModelTester(self   )
							lowercase_ : Optional[Any]							       =   ConfigTester(self    ,						config_class=A    ,						emb_dim=37   )




					def   A  (				self     :						List[str]   )  ->  Dict:
							self.config_tester.run_common_tests()




					def   A  (				self     :						List[str]   )  ->  int:
							lowercase_ : Tuple							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_model(*A   )




					def   A  (				self     :						List[str]   )  ->  int:
							lowercase_ : str							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_lm_head(*A   )




					def   A  (				self     :						List[str]   )  ->  int:
							lowercase_ : Tuple							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_qa(*A   )




					def   A  (				self     :						Any   )  ->  List[Any]:
							lowercase_ : Any							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_sequence_classif(*A   )




					def   A  (				self     :						Optional[int]   )  ->  str:
							lowercase_ : Union[str, Any]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_for_token_classification(*A   )




					def   A  (				self     :						Optional[Any]   )  ->  Tuple:
							lowercase_ : Tuple							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_for_multiple_choice(*A   )






					@slow
					def   A  (				self     :						List[Any]   )  ->  Optional[Any]:
							for model_name in TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowercase_ : Tuple							       =   TFFlaubertModel.from_pretrained(A   )
									self.assertIsNotNone(A   )





@require_tf
@require_sentencepiece
@require_tokenizers
class 					_UpperCAmelCase       ( unittest.TestCase							):




					@slow
					def   A  (				self     :						str   )  ->  int:
							lowercase_ : Any							       =   TFFlaubertModel.from_pretrained('''jplu/tf-flaubert-small-cased'''   )

							lowercase_ : Optional[int]							       =   tf.convert_to_tensor(
							    [[0, 1_58, 7_35, 25_92, 14_24, 67_27, 82, 1]]    ,						dtype=tf.intaa    ,						)  # "J'aime flaubert !"

							lowercase_ : int							       =   model(A   )[0]
							lowercase_ : str							       =   tf.TensorShape((1, 8, 5_12)   )
							self.assertEqual(output.shape    ,						A   )
							# compare the actual values for a slice.
							lowercase_ : int							       =   tf.convert_to_tensor(
							    [
							        [
							            [-1.8768773, -1.566555, 0.27072418],
							            [-1.6920038, -0.5873505, 1.9329599],
							            [-2.9563985, -1.6993835, 1.7972052],
							        ]
							    ]    ,						dtype=tf.floataa    ,						)

							self.assertTrue(np.allclose(output[:, :3, :3].numpy()    ,						expected_slice.numpy()    ,						atol=1e-4   )   )










import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs

@pytest.mark.parametrize(
    'kwargs, expected',	[
        ({'num_shards': 0, 'max_num_jobs': 1}, []),
        ({'num_shards': 10, 'max_num_jobs': 1}, [range(10	)]),
        ({'num_shards': 10, 'max_num_jobs': 10}, [range(lowerCAmelCase_,	i + 1	) for i in range(10	)]),
        ({'num_shards': 1, 'max_num_jobs': 10}, [range(1	)]),
        ({'num_shards': 10, 'max_num_jobs': 3}, [range(0,	4	), range(4,	7	), range(7,	10	)]),
        ({'num_shards': 3, 'max_num_jobs': 10}, [range(0,	1	), range(1,	2	), range(2,	3	)]),
    ],	)
def   a_    ( lowerCAmelCase_      :			Optional[Any],	lowerCAmelCase_      :			Any	):
   __lowerCAmelCase								=	_distribute_shards(**lowerCAmelCase_	)
   assert out == expected

@pytest.mark.parametrize(
    'gen_kwargs, max_num_jobs, expected',	[
        ({'foo': 0}, 10, [{'foo': 0}]),
        ({'shards': [0, 1, 2, 3]}, 1, [{'shards': [0, 1, 2, 3]}]),
        ({'shards': [0, 1, 2, 3]}, 4, [{'shards': [0]}, {'shards': [1]}, {'shards': [2]}, {'shards': [3]}]),
        ({'shards': [0, 1]}, 4, [{'shards': [0]}, {'shards': [1]}]),
        ({'shards': [0, 1, 2, 3]}, 2, [{'shards': [0, 1]}, {'shards': [2, 3]}]),
    ],	)
def   a_    ( lowerCAmelCase_      :			Optional[Any],	lowerCAmelCase_      :			List[str],	lowerCAmelCase_      :			Optional[int]	):
   __lowerCAmelCase								=	_split_gen_kwargs(lowerCAmelCase_,	lowerCAmelCase_	)
   assert out == expected




@pytest.mark.parametrize(
    'gen_kwargs, expected',	[
        ({'foo': 0}, 1),
        ({'shards': [0]}, 1),
        ({'shards': [0, 1, 2, 3]}, 4),
        ({'shards': [0, 1, 2, 3], 'foo': 0}, 4),
        ({'shards': [0, 1, 2, 3], 'other': (0, 1)}, 4),
        ({'shards': [0, 1, 2, 3], 'shards2': [0, 1]}, RuntimeError),
    ],	)
def   a_    ( lowerCAmelCase_      :			int,	lowerCAmelCase_      :			Any	):
   if expected is RuntimeError:
      with pytest.raises(lowerCAmelCase_	):
         _number_of_shards_in_gen_kwargs(lowerCAmelCase_	)
   else:
      __lowerCAmelCase								=	_number_of_shards_in_gen_kwargs(lowerCAmelCase_	)
      assert out == expected

