


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






			"""simple docstring"""






			def __init__(				self   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							parent
						UpperCamelCase__       =							13
						UpperCamelCase__       =							7
						UpperCamelCase__       =							True
						UpperCamelCase__       =							True
						UpperCamelCase__       =							True
						UpperCamelCase__       =							True
						UpperCamelCase__       =							True
						UpperCamelCase__       =							False
						UpperCamelCase__       =							False
						UpperCamelCase__       =							False
						UpperCamelCase__       =							2
						UpperCamelCase__       =							99
						UpperCamelCase__       =							0
						UpperCamelCase__       =							32
						UpperCamelCase__       =							2
						UpperCamelCase__       =							4
						UpperCamelCase__       =							0.1
						UpperCamelCase__       =							0.1
						UpperCamelCase__       =							512
						UpperCamelCase__       =							16
						UpperCamelCase__       =							2
						UpperCamelCase__       =							0.02
						UpperCamelCase__       =							3
						UpperCamelCase__       =							4
						UpperCamelCase__       =							"last"
						UpperCamelCase__       =							True
						UpperCamelCase__       =							None
						UpperCamelCase__       =							0




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size  )
						UpperCamelCase__       =							random_attention_mask([self.batch_size, self.seq_length]   ,    dtype=tf.floataa  )

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

						UpperCamelCase__       =							None
						if self.use_token_type_ids:
									UpperCamelCase__       =							ids_tensor([self.batch_size, self.seq_length]   ,    self.n_langs  )

						UpperCamelCase__       =							None
						UpperCamelCase__       =							None
						UpperCamelCase__       =							None
						if self.use_labels:
									UpperCamelCase__       =							ids_tensor([self.batch_size]   ,    self.type_sequence_label_size  )
									UpperCamelCase__       =							ids_tensor([self.batch_size, self.seq_length]   ,    self.num_labels  )
									UpperCamelCase__       =							ids_tensor([self.batch_size]   ,    2   ,    dtype=tf.floataa  )
									UpperCamelCase__       =							ids_tensor([self.batch_size]   ,    self.num_choices  )

						UpperCamelCase__       =							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 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							TFFlaubertModel(config=snake_case  )
						UpperCamelCase__       =							{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
						UpperCamelCase__       =							model(snake_case  )

						UpperCamelCase__       =							[input_ids, input_mask]
						UpperCamelCase__       =							model(snake_case  )
						self.parent.assertEqual(result.last_hidden_state.shape   ,    (self.batch_size, self.seq_length, self.hidden_size)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							TFFlaubertWithLMHeadModel(snake_case  )

						UpperCamelCase__       =							{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
						UpperCamelCase__       =							model(snake_case  )

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




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							TFFlaubertForQuestionAnsweringSimple(snake_case  )

						UpperCamelCase__       =							{"input_ids": input_ids, "lengths": input_lengths}

						UpperCamelCase__       =							model(snake_case  )

						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 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							TFFlaubertForSequenceClassification(snake_case  )

						UpperCamelCase__       =							{"input_ids": input_ids, "lengths": input_lengths}

						UpperCamelCase__       =							model(snake_case  )

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




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							self.num_labels
						UpperCamelCase__       =							TFFlaubertForTokenClassification(config=snake_case  )
						UpperCamelCase__       =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						UpperCamelCase__       =							model(snake_case  )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.num_labels)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							self.num_choices
						UpperCamelCase__       =							TFFlaubertForMultipleChoice(config=snake_case  )
						UpperCamelCase__       =							tf.tile(tf.expand_dims(snake_case   ,    1  )   ,    (1, self.num_choices, 1)  )
						UpperCamelCase__       =							tf.tile(tf.expand_dims(snake_case   ,    1  )   ,    (1, self.num_choices, 1)  )
						UpperCamelCase__       =							tf.tile(tf.expand_dims(snake_case   ,    1  )   ,    (1, self.num_choices, 1)  )
						UpperCamelCase__       =							{
						    "input_ids": multiple_choice_inputs_ids,
						    "attention_mask": multiple_choice_input_mask,
						    "token_type_ids": multiple_choice_token_type_ids,
						}
						UpperCamelCase__       =							model(snake_case  )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.num_choices)  )







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.prepare_config_and_inputs()
						(
						    (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   
						)       =							config_and_inputs
						UpperCamelCase__       =							{
						    "input_ids": input_ids,
						    "token_type_ids": token_type_ids,
						    "langs": token_type_ids,
						    "lengths": input_lengths,
						}
						return config, inputs_dict







@require_tf
class 							lowerCamelCase__						(				UpperCAmelCase       ,						UpperCAmelCase       ,						unittest.TestCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       Optional[int] 			=		(
			    (
			        TFFlaubertModel,
			        TFFlaubertWithLMHeadModel,
			        TFFlaubertForSequenceClassification,
			        TFFlaubertForQuestionAnsweringSimple,
			        TFFlaubertForTokenClassification,
			        TFFlaubertForMultipleChoice,
			    )
			    if is_tf_available()
			    else ()
			)
			_UpperCamelCase		:       Any 			=		(
			    (TFFlaubertWithLMHeadModel,) if is_tf_available() else ()
			)  # TODO (PVP): Check other models whether language generation is also applicable
			_UpperCamelCase		:       str 			=		(
			    {
			        'feature-extraction': TFFlaubertModel,
			        'fill-mask': TFFlaubertWithLMHeadModel,
			        'question-answering': TFFlaubertForQuestionAnsweringSimple,
			        'text-classification': TFFlaubertForSequenceClassification,
			        'token-classification': TFFlaubertForTokenClassification,
			        'zero-shot': TFFlaubertForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			_UpperCamelCase		:       str 			=		False
			_UpperCamelCase		:       Dict 			=		False




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						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 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							TFFlaubertModelTester(self  )
						UpperCamelCase__       =							ConfigTester(self   ,    config_class=snake_case   ,    emb_dim=37  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						self.config_tester.run_common_tests()




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_flaubert_model(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_flaubert_lm_head(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_flaubert_qa(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_flaubert_sequence_classif(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_flaubert_for_token_classification(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_flaubert_for_multiple_choice(*snake_case  )







			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for model_name in TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCamelCase__       =							TFFlaubertModel.from_pretrained(snake_case  )
									self.assertIsNotNone(snake_case  )







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






			"""simple docstring"""






			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							TFFlaubertModel.from_pretrained("jplu/tf-flaubert-small-cased"  )

						UpperCamelCase__       =							tf.convert_to_tensor(
						    [[0, 158, 735, 2592, 1424, 6727, 82, 1]]   ,    dtype=tf.intaa   ,    )  # "J'aime flaubert !"

						UpperCamelCase__       =							model(snake_case  )[0]
						UpperCamelCase__       =							tf.TensorShape((1, 8, 512)  )
						self.assertEqual(output.shape   ,    snake_case  )
						# compare the actual values for a slice.
						UpperCamelCase__       =							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 torch

from diffusers import DDIMParallelScheduler

from .test_schedulers import SchedulerCommonTest







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       int 			=		(DDIMParallelScheduler,)
			_UpperCamelCase		:       List[Any] 			=		(('eta', 0.0), ('num_inference_steps', 50))




			def 					snake_case__					(				self   ,    **snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							{
						    "num_train_timesteps": 1000,
						    "beta_start": 0.0001,
						    "beta_end": 0.02,
						    "beta_schedule": "linear",
						    "clip_sample": True,
						}

						config.update(**snake_case  )
						return config




			def 					snake_case__					(				self   ,    **snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.scheduler_classes[0]
						UpperCamelCase__       =							self.get_scheduler_config(**snake_case  )
						UpperCamelCase__       =							scheduler_class(**snake_case  )

						UpperCamelCase__,   UpperCamelCase__       =							10, 0.0

						UpperCamelCase__       =							self.dummy_model()
						UpperCamelCase__       =							self.dummy_sample_deter

						scheduler.set_timesteps(snake_case  )

						for t in scheduler.timesteps:
									UpperCamelCase__       =							model(snake_case   ,    snake_case  )
									UpperCamelCase__       =							scheduler.step(snake_case   ,    snake_case   ,    snake_case   ,    snake_case  ).prev_sample

						return sample




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for timesteps in [100, 500, 1000]:
									self.check_over_configs(num_train_timesteps=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for steps_offset in [0, 1]:
									self.check_over_configs(steps_offset=snake_case  )

						UpperCamelCase__       =							self.scheduler_classes[0]
						UpperCamelCase__       =							self.get_scheduler_config(steps_offset=1  )
						UpperCamelCase__       =							scheduler_class(**snake_case  )
						scheduler.set_timesteps(5  )
						assert torch.equal(scheduler.timesteps   ,    torch.LongTensor([801, 601, 401, 201, 1]  )  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1]   ,    [0.002, 0.02, 0.2, 2]  ):
									self.check_over_configs(beta_start=snake_case   ,    beta_end=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for schedule in ["linear", "squaredcos_cap_v2"]:
									self.check_over_configs(beta_schedule=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for prediction_type in ["epsilon", "v_prediction"]:
									self.check_over_configs(prediction_type=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for clip_sample in [True, False]:
									self.check_over_configs(clip_sample=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for timestep_spacing in ["trailing", "leading"]:
									self.check_over_configs(timestep_spacing=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for rescale_betas_zero_snr in [True, False]:
									self.check_over_configs(rescale_betas_zero_snr=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						self.check_over_configs(thresholding=snake_case  )
						for threshold in [0.5, 1.0, 2.0]:
									for prediction_type in ["epsilon", "v_prediction"]:
												self.check_over_configs(
												    thresholding=snake_case   ,    prediction_type=snake_case   ,    sample_max_value=snake_case   ,    )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for t in [1, 10, 49]:
									self.check_over_forward(time_step=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for t, num_inference_steps in zip([1, 10, 50]   ,    [10, 50, 500]  ):
									self.check_over_forward(time_step=snake_case   ,    num_inference_steps=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for t, eta in zip([1, 10, 49]   ,    [0.0, 0.5, 1.0]  ):
									self.check_over_forward(time_step=snake_case   ,    eta=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.scheduler_classes[0]
						UpperCamelCase__       =							self.get_scheduler_config()
						UpperCamelCase__       =							scheduler_class(**snake_case  )

						assert torch.sum(torch.abs(scheduler._get_variance(0   ,    0  ) - 0.0  )  ) < 1E-5
						assert torch.sum(torch.abs(scheduler._get_variance(420   ,    400  ) - 0.14771  )  ) < 1E-5
						assert torch.sum(torch.abs(scheduler._get_variance(980   ,    960  ) - 0.32460  )  ) < 1E-5
						assert torch.sum(torch.abs(scheduler._get_variance(0   ,    0  ) - 0.0  )  ) < 1E-5
						assert torch.sum(torch.abs(scheduler._get_variance(487   ,    486  ) - 0.00979  )  ) < 1E-5
						assert torch.sum(torch.abs(scheduler._get_variance(999   ,    998  ) - 0.02  )  ) < 1E-5




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.scheduler_classes[0]
						UpperCamelCase__       =							self.get_scheduler_config()
						UpperCamelCase__       =							scheduler_class(**snake_case  )

						UpperCamelCase__,   UpperCamelCase__       =							10, 0.0
						scheduler.set_timesteps(snake_case  )

						UpperCamelCase__       =							self.dummy_model()
						UpperCamelCase__       =							self.dummy_sample_deter
						UpperCamelCase__       =							self.dummy_sample_deter + 0.1
						UpperCamelCase__       =							self.dummy_sample_deter - 0.1

						UpperCamelCase__       =							samplea.shape[0]
						UpperCamelCase__       =							torch.stack([samplea, samplea, samplea]   ,    dim=0  )
						UpperCamelCase__       =							torch.arange(snake_case  )[0:3, None].repeat(1   ,    snake_case  )

						UpperCamelCase__       =							model(samples.flatten(0   ,    1  )   ,    timesteps.flatten(0   ,    1  )  )
						UpperCamelCase__       =							scheduler.batch_step_no_noise(snake_case   ,    timesteps.flatten(0   ,    1  )   ,    samples.flatten(0   ,    1  )   ,    snake_case  )

						UpperCamelCase__       =							torch.sum(torch.abs(snake_case  )  )
						UpperCamelCase__       =							torch.mean(torch.abs(snake_case  )  )

						assert abs(result_sum.item() - 1147.7904  ) < 1E-2
						assert abs(result_mean.item() - 0.4982  ) < 1E-3




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.full_loop()

						UpperCamelCase__       =							torch.sum(torch.abs(snake_case  )  )
						UpperCamelCase__       =							torch.mean(torch.abs(snake_case  )  )

						assert abs(result_sum.item() - 172.0067  ) < 1E-2
						assert abs(result_mean.item() - 0.223967  ) < 1E-3




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.full_loop(prediction_type="v_prediction"  )

						UpperCamelCase__       =							torch.sum(torch.abs(snake_case  )  )
						UpperCamelCase__       =							torch.mean(torch.abs(snake_case  )  )

						assert abs(result_sum.item() - 52.5302  ) < 1E-2
						assert abs(result_mean.item() - 0.0684  ) < 1E-3




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.full_loop(set_alpha_to_one=snake_case   ,    beta_start=0.01  )
						UpperCamelCase__       =							torch.sum(torch.abs(snake_case  )  )
						UpperCamelCase__       =							torch.mean(torch.abs(snake_case  )  )

						assert abs(result_sum.item() - 149.8295  ) < 1E-2
						assert abs(result_mean.item() - 0.1951  ) < 1E-3







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.full_loop(set_alpha_to_one=snake_case   ,    beta_start=0.01  )
						UpperCamelCase__       =							torch.sum(torch.abs(snake_case  )  )
						UpperCamelCase__       =							torch.mean(torch.abs(snake_case  )  )

						assert abs(result_sum.item() - 149.0784  ) < 1E-2
						assert abs(result_mean.item() - 0.1941  ) < 1E-3


