

import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    FlaubertForMultipleChoice,
					    FlaubertForQuestionAnswering,
					    FlaubertForQuestionAnsweringSimple,
					    FlaubertForSequenceClassification,
					    FlaubertForTokenClassification,
					    FlaubertModel,
					    FlaubertWithLMHeadModel,
					)
					from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST




class     a			(				SCREAMING_SNAKE_CASE       ):






						"""simple docstring"""




						def __init__(							self					:      Optional[Any]			,			snake_case_					:      Optional[Any]			,			snake_case_					:      Tuple=1_3			,			snake_case_					:      Union[str, Any]=7			,			snake_case_					:      Union[str, Any]=True			,			snake_case_					:      Any=True			,			snake_case_					:      Dict=True			,			snake_case_					:      Tuple=True			,			snake_case_					:      Dict=True			,			snake_case_					:      Dict=False			,			snake_case_					:      Union[str, Any]=False			,			snake_case_					:      Optional[int]=False			,			snake_case_					:      Optional[int]=2			,			snake_case_					:      List[str]=9_9			,			snake_case_					:      List[str]=0			,			snake_case_					:      Dict=3_2			,			snake_case_					:      int=5			,			snake_case_					:      int=4			,			snake_case_					:      Tuple=0.1			,			snake_case_					:      int=0.1			,			snake_case_					:      Any=5_1_2			,			snake_case_					:      Optional[Any]=1_2			,			snake_case_					:      List[str]=2			,			snake_case_					:      Union[str, Any]=0.0_2			,			snake_case_					:      Any=3			,			snake_case_					:      str=4			,			snake_case_					:      Optional[Any]="last"			,			snake_case_					:      Optional[int]=None			,			snake_case_					:      List[Any]=None			,			):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  parent
								snake_case__						:       str       						=  batch_size
								snake_case__						:       Union[str, Any]       						=  seq_length
								snake_case__						:       List[str]       						=  is_training
								snake_case__						:       Optional[Any]       						=  use_input_lengths
								snake_case__						:       List[str]       						=  use_token_type_ids
								snake_case__						:       Dict       						=  use_labels
								snake_case__						:       List[str]       						=  gelu_activation
								snake_case__						:       Optional[Any]       						=  sinusoidal_embeddings
								snake_case__						:       Any       						=  causal
								snake_case__						:       Dict       						=  asm
								snake_case__						:       Union[str, Any]       						=  n_langs
								snake_case__						:       List[str]       						=  vocab_size
								snake_case__						:       Optional[int]       						=  n_special
								snake_case__						:       List[Any]       						=  hidden_size
								snake_case__						:       Dict       						=  num_hidden_layers
								snake_case__						:       Dict       						=  num_attention_heads
								snake_case__						:       List[Any]       						=  hidden_dropout_prob
								snake_case__						:       Tuple       						=  attention_probs_dropout_prob
								snake_case__						:       int       						=  max_position_embeddings
								snake_case__						:       Dict       						=  type_vocab_size
								snake_case__						:       Tuple       						=  type_sequence_label_size
								snake_case__						:       Dict       						=  initializer_range
								snake_case__						:       List[Any]       						=  num_labels
								snake_case__						:       Optional[int]       						=  num_choices
								snake_case__						:       int       						=  summary_type
								snake_case__						:       Union[str, Any]       						=  use_proj
								snake_case__						:       int       						=  scope







						def       __magic_name__							(							self					:      Any  ):



								'''simple docstring'''



								snake_case__						:       Union[str, Any]       						=  ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size  )
								snake_case__						:       Dict       						=  random_attention_mask([self.batch_size, self.seq_length]  )

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

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

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

								snake_case__						:       List[Any]       						=  self.get_config()

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







						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								return 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			,			)







						def       __magic_name__							(							self					:      Union[str, Any]			,			snake_case_					:      List[Any]			,			snake_case_					:      str			,			snake_case_					:      List[str]			,			snake_case_					:      Tuple			,			snake_case_					:      Optional[int]			,			snake_case_					:      Optional[int]			,			snake_case_					:      List[str]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      str			,			):



								'''simple docstring'''



								snake_case__						:       Union[str, Any]       						=  FlaubertModel(config=snake_case_  )
								model.to(snake_case_  )
								model.eval()
								snake_case__						:       int       						=  model(snake_case_			,			lengths=snake_case_			,			langs=snake_case_  )
								snake_case__						:       Optional[int]       						=  model(snake_case_			,			langs=snake_case_  )
								snake_case__						:       Any       						=  model(snake_case_  )
								self.parent.assertEqual(result.last_hidden_state.shape			,			(self.batch_size, self.seq_length, self.hidden_size)  )







						def       __magic_name__							(							self					:      List[str]			,			snake_case_					:      Dict			,			snake_case_					:      List[Any]			,			snake_case_					:      List[str]			,			snake_case_					:      List[str]			,			snake_case_					:      List[str]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      List[str]			,			snake_case_					:      List[Any]			,			snake_case_					:      int			,			):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  FlaubertWithLMHeadModel(snake_case_  )
								model.to(snake_case_  )
								model.eval()

								snake_case__						:       Any       						=  model(snake_case_			,			token_type_ids=snake_case_			,			labels=snake_case_  )
								self.parent.assertEqual(result.loss.shape			,			()  )
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.vocab_size)  )







						def       __magic_name__							(							self					:      Optional[Any]			,			snake_case_					:      Tuple			,			snake_case_					:      int			,			snake_case_					:      Tuple			,			snake_case_					:      Dict			,			snake_case_					:      Optional[int]			,			snake_case_					:      int			,			snake_case_					:      List[str]			,			snake_case_					:      Optional[Any]			,			snake_case_					:      Dict			,			):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  FlaubertForQuestionAnsweringSimple(snake_case_  )
								model.to(snake_case_  )
								model.eval()

								snake_case__						:       List[Any]       						=  model(snake_case_  )

								snake_case__						:       List[str]       						=  model(snake_case_			,			start_positions=snake_case_			,			end_positions=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       __magic_name__							(							self					:      Optional[int]			,			snake_case_					:      Any			,			snake_case_					:      List[Any]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      str			,			snake_case_					:      str			,			snake_case_					:      int			,			snake_case_					:      Optional[Any]			,			snake_case_					:      List[str]			,			snake_case_					:      Tuple			,			):



								'''simple docstring'''



								snake_case__						:       Tuple       						=  FlaubertForQuestionAnswering(snake_case_  )
								model.to(snake_case_  )
								model.eval()

								snake_case__						:       Optional[int]       						=  model(snake_case_  )

								snake_case__						:       int       						=  model(
								    snake_case_			,			start_positions=snake_case_			,			end_positions=snake_case_			,			cls_index=snake_case_			,			is_impossible=snake_case_			,			p_mask=snake_case_			,			)

								snake_case__						:       Any       						=  model(
								    snake_case_			,			start_positions=snake_case_			,			end_positions=snake_case_			,			cls_index=snake_case_			,			is_impossible=snake_case_			,			)

								(snake_case__       )						:       List[str]       						=  result_with_labels.to_tuple()

								snake_case__						:       str       						=  model(snake_case_			,			start_positions=snake_case_			,			end_positions=snake_case_  )

								(snake_case__       )						:       Tuple       						=  result_with_labels.to_tuple()

								self.parent.assertEqual(result_with_labels.loss.shape			,			()  )
								self.parent.assertEqual(result.start_top_log_probs.shape			,			(self.batch_size, model.config.start_n_top)  )
								self.parent.assertEqual(result.start_top_index.shape			,			(self.batch_size, model.config.start_n_top)  )
								self.parent.assertEqual(
								    result.end_top_log_probs.shape			,			(self.batch_size, model.config.start_n_top * model.config.end_n_top)  )
								self.parent.assertEqual(
								    result.end_top_index.shape			,			(self.batch_size, model.config.start_n_top * model.config.end_n_top)  )
								self.parent.assertEqual(result.cls_logits.shape			,			(self.batch_size,)  )







						def       __magic_name__							(							self					:      Union[str, Any]			,			snake_case_					:      Any			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      Any			,			snake_case_					:      Any			,			snake_case_					:      Any			,			snake_case_					:      int			,			snake_case_					:      List[str]			,			snake_case_					:      List[str]			,			snake_case_					:      List[str]			,			):



								'''simple docstring'''



								snake_case__						:       str       						=  FlaubertForSequenceClassification(snake_case_  )
								model.to(snake_case_  )
								model.eval()

								snake_case__						:       Dict       						=  model(snake_case_  )
								snake_case__						:       str       						=  model(snake_case_			,			labels=snake_case_  )

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







						def       __magic_name__							(							self					:      Dict			,			snake_case_					:      Tuple			,			snake_case_					:      Any			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      Any			,			snake_case_					:      str			,			snake_case_					:      int			,			snake_case_					:      Union[str, Any]			,			):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  self.num_labels
								snake_case__						:       int       						=  FlaubertForTokenClassification(snake_case_  )
								model.to(snake_case_  )
								model.eval()

								snake_case__						:       Union[str, Any]       						=  model(snake_case_			,			attention_mask=snake_case_			,			labels=snake_case_  )
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.num_labels)  )







						def       __magic_name__							(							self					:      Optional[Any]			,			snake_case_					:      Optional[int]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      int			,			snake_case_					:      int			,			snake_case_					:      Tuple			,			snake_case_					:      List[str]			,			snake_case_					:      Optional[Any]			,			snake_case_					:      str			,			snake_case_					:      Union[str, Any]			,			):



								'''simple docstring'''



								snake_case__						:       Union[str, Any]       						=  self.num_choices
								snake_case__						:       List[Any]       						=  FlaubertForMultipleChoice(config=snake_case_  )
								model.to(snake_case_  )
								model.eval()
								snake_case__						:       Optional[int]       						=  input_ids.unsqueeze(1  ).expand(-1			,			self.num_choices			,			-1  ).contiguous()
								snake_case__						:       Union[str, Any]       						=  token_type_ids.unsqueeze(1  ).expand(-1			,			self.num_choices			,			-1  ).contiguous()
								snake_case__						:       Any       						=  input_mask.unsqueeze(1  ).expand(-1			,			self.num_choices			,			-1  ).contiguous()
								snake_case__						:       List[str]       						=  model(
								    snake_case_			,			attention_mask=snake_case_			,			token_type_ids=snake_case_			,			labels=snake_case_			,			)
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.num_choices)  )







						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						:       int       						=  self.prepare_config_and_inputs()
								(
								    snake_case__       
								)						:       Tuple       						=  config_and_inputs
								snake_case__						:       Optional[Any]       						=  {
								    '''input_ids''': input_ids,
								    '''token_type_ids''': token_type_ids,
								    '''lengths''': input_lengths,
								    '''attention_mask''': input_mask,
								}
								return config, inputs_dict





@require_torch
class     a			(				SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    (
						    (
						        FlaubertModel,
						        FlaubertWithLMHeadModel,
						        FlaubertForQuestionAnswering,
						        FlaubertForQuestionAnsweringSimple,
						        FlaubertForSequenceClassification,
						        FlaubertForTokenClassification,
						        FlaubertForMultipleChoice,
						    )
						    if is_torch_available()
						    else ()
						)
						__UpperCAmelCase			     =    (
						    {
						        """feature-extraction""": FlaubertModel,
						        """fill-mask""": FlaubertWithLMHeadModel,
						        """question-answering""": FlaubertForQuestionAnsweringSimple,
						        """text-classification""": FlaubertForSequenceClassification,
						        """token-classification""": FlaubertForTokenClassification,
						        """zero-shot""": FlaubertForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)







						def       __magic_name__							(							self					:      List[str]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      List[str]			,			snake_case_					:      List[str]			,			snake_case_					:      Tuple  ):



								'''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       __magic_name__							(							self					:      Dict			,			snake_case_					:      str			,			snake_case_					:      int			,			snake_case_					:      Any=False  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  super()._prepare_for_class(snake_case_			,			snake_case_			,			return_labels=snake_case_  )

								if return_labels:
										if model_class.__name__ == "FlaubertForQuestionAnswering":
												snake_case__						:       str       						=  torch.zeros(
												    self.model_tester.batch_size			,			dtype=torch.long			,			device=snake_case_  )
												snake_case__						:       Union[str, Any]       						=  torch.zeros(
												    self.model_tester.batch_size			,			dtype=torch.long			,			device=snake_case_  )

								return inputs_dict







						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								snake_case__						:       Dict       						=  FlaubertModelTester(self  )
								snake_case__						:       str       						=  ConfigTester(self			,			config_class=snake_case_			,			emb_dim=3_7  )







						def       __magic_name__							(							self					:      Optional[Any]  ):



								'''simple docstring'''



								self.config_tester.run_common_tests()







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       str       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_model(*snake_case_  )







						def       __magic_name__							(							self					:      Tuple  ):



								'''simple docstring'''



								snake_case__						:       Tuple       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_lm_head(*snake_case_  )







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       Optional[Any]       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_simple_qa(*snake_case_  )







						def       __magic_name__							(							self					:      List[str]  ):



								'''simple docstring'''



								snake_case__						:       Tuple       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_qa(*snake_case_  )







						def       __magic_name__							(							self					:      Tuple  ):



								'''simple docstring'''



								snake_case__						:       Optional[Any]       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_sequence_classif(*snake_case_  )







						def       __magic_name__							(							self					:      Tuple  ):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_token_classif(*snake_case_  )







						def       __magic_name__							(							self					:      Tuple  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_multiple_choice(*snake_case_  )







						@slow
						def       __magic_name__							(							self					:      Optional[Any]  ):



								'''simple docstring'''



								for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										snake_case__						:       Any       						=  FlaubertModel.from_pretrained(snake_case_  )
										self.assertIsNotNone(snake_case_  )







						@slow
						@require_torch_gpu
						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
										# FlauBertForMultipleChoice behaves incorrectly in JIT environments.
										if model_class == FlaubertForMultipleChoice:
												return

										snake_case__						:       Tuple       						=  True
										snake_case__						:       Tuple       						=  model_class(config=snake_case_  )

										snake_case__						:       Optional[int]       						=  self._prepare_for_class(snake_case_			,			snake_case_  )
										snake_case__						:       List[str]       						=  torch.jit.trace(
										    snake_case_			,			(inputs_dict['''input_ids'''].to('''cpu'''  ), inputs_dict['''attention_mask'''].to('''cpu'''  ))  )

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





@require_torch
class     a			(				unittest.TestCase       ):






						"""simple docstring"""




						@slow
						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						:       List[Any]       						=  FlaubertModel.from_pretrained('''flaubert/flaubert_base_cased'''  )
								snake_case__						:       Dict       						=  torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]  )
								with torch.no_grad():
										snake_case__						:       Any       						=  model(snake_case_  )[0]
								snake_case__						:       List[str]       						=  torch.Size((1, 1_1, 7_6_8)  )
								self.assertEqual(output.shape			,			snake_case_  )
								snake_case__						:       Any       						=  torch.tensor(
								    [[[-2.6_2_5_1, -1.4_2_9_8, -0.0_2_2_7], [-2.8_5_1_0, -1.6_3_8_7, 0.2_2_5_8], [-2.8_1_1_4, -1.1_8_3_2, -0.3_0_6_6]]]  )

								self.assertTrue(torch.allclose(output[:, :3, :3]			,			snake_case_			,			atol=1e-4  )  )






'''simple docstring'''


import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
					import jax




@require_flax
class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    FlaxAutoencoderKL







						@property
						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  4
								snake_case__						:       Optional[Any]       						=  3
								snake_case__						:       Optional[int]       						=  (3_2, 3_2)

								snake_case__						:       Optional[int]       						=  jax.random.PRNGKey(0  )
								snake_case__						:       Union[str, Any]       						=  jax.random.uniform(snake_case_			,			((batch_size, num_channels) + sizes)  )

								return {"sample": image, "prng_key": prng_key}







						def       __magic_name__							(							self					:      Union[str, Any]  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  {
								    '''block_out_channels''': [3_2, 6_4],
								    '''in_channels''': 3,
								    '''out_channels''': 3,
								    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
								    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
								    '''latent_channels''': 4,
								}
								snake_case__						:       Dict       						=  self.dummy_input
								return init_dict, inputs_dict



