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        lowerCamelCase__				(   UpperCAmelCase				):





			def __init__(       self	:	List[Any]  ,  lowercase__	:	str  ,  lowercase__	:	List[Any]=13  ,  lowercase__	:	Optional[Any]=7  ,  lowercase__	:	Optional[int]=True  ,  lowercase__	:	Optional[int]=True  ,  lowercase__	:	Optional[Any]=True  ,  lowercase__	:	List[str]=True  ,  lowercase__	:	Union[str, Any]=True  ,  lowercase__	:	Any=False  ,  lowercase__	:	Optional[int]=False  ,  lowercase__	:	List[Any]=False  ,  lowercase__	:	int=2  ,  lowercase__	:	str=99  ,  lowercase__	:	Optional[Any]=0  ,  lowercase__	:	Optional[int]=32  ,  lowercase__	:	Tuple=5  ,  lowercase__	:	Union[str, Any]=4  ,  lowercase__	:	Optional[Any]=0.1  ,  lowercase__	:	Dict=0.1  ,  lowercase__	:	List[Any]=5_12  ,  lowercase__	:	str=12  ,  lowercase__	:	Optional[int]=2  ,  lowercase__	:	Dict=0.0_2  ,  lowercase__	:	int=3  ,  lowercase__	:	str=4  ,  lowercase__	:	Union[str, Any]="last"  ,  lowercase__	:	List[str]=None  ,  lowercase__	:	Optional[Any]=None  ,  ):
										_lowerCAmelCase							=							parent
										_lowerCAmelCase							=							batch_size
										_lowerCAmelCase							=							seq_length
										_lowerCAmelCase							=							is_training
										_lowerCAmelCase							=							use_input_lengths
										_lowerCAmelCase							=							use_token_type_ids
										_lowerCAmelCase							=							use_labels
										_lowerCAmelCase							=							gelu_activation
										_lowerCAmelCase							=							sinusoidal_embeddings
										_lowerCAmelCase							=							causal
										_lowerCAmelCase							=							asm
										_lowerCAmelCase							=							n_langs
										_lowerCAmelCase							=							vocab_size
										_lowerCAmelCase							=							n_special
										_lowerCAmelCase							=							hidden_size
										_lowerCAmelCase							=							num_hidden_layers
										_lowerCAmelCase							=							num_attention_heads
										_lowerCAmelCase							=							hidden_dropout_prob
										_lowerCAmelCase							=							attention_probs_dropout_prob
										_lowerCAmelCase							=							max_position_embeddings
										_lowerCAmelCase							=							type_vocab_size
										_lowerCAmelCase							=							type_sequence_label_size
										_lowerCAmelCase							=							initializer_range
										_lowerCAmelCase							=							num_labels
										_lowerCAmelCase							=							num_choices
										_lowerCAmelCase							=							summary_type
										_lowerCAmelCase							=							use_proj
										_lowerCAmelCase							=							scope





			def   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										_lowerCAmelCase							=							ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size						)
										_lowerCAmelCase							=							random_attention_mask([self.batch_size, self.seq_length]						)

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

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

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

										_lowerCAmelCase							=							self.get_config()

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										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   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	str  ,  lowercase__	:	Any  ,  lowercase__	:	List[Any]  ,  lowercase__	:	str  ,  lowercase__	:	Dict  ,  lowercase__	:	int  ,  lowercase__	:	Tuple  ,  lowercase__	:	Dict  ,  ):
										_lowerCAmelCase							=							FlaubertModel(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  lengths=lowercase__  ,  langs=lowercase__						)
										_lowerCAmelCase							=							model(lowercase__  ,  langs=lowercase__						)
										_lowerCAmelCase							=							model(lowercase__						)
										self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	int  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int  ,  lowercase__	:	List[str]  ,  lowercase__	:	List[str]  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	List[str]  ,  lowercase__	:	int  ,  ):
										_lowerCAmelCase							=							FlaubertWithLMHeadModel(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__  ,  token_type_ids=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.loss.shape  ,  ()						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]  ,  lowercase__	:	int  ,  lowercase__	:	List[str]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	int  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Dict  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int  ,  lowercase__	:	Optional[int]  ,  ):
										_lowerCAmelCase							=							FlaubertForQuestionAnsweringSimple(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__						)

										_lowerCAmelCase							=							model(lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__						)
										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	:	Union[str, Any]  ,  lowercase__	:	Any  ,  lowercase__	:	int  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	str  ,  lowercase__	:	str  ,  lowercase__	:	Tuple  ,  lowercase__	:	int  ,  lowercase__	:	List[str]  ,  ):
										_lowerCAmelCase							=							FlaubertForQuestionAnswering(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__						)

										_lowerCAmelCase							=							model(
										    lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__  ,  cls_index=lowercase__  ,  is_impossible=lowercase__  ,  p_mask=lowercase__  ,  )

										_lowerCAmelCase							=							model(
										    lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__  ,  cls_index=lowercase__  ,  is_impossible=lowercase__  ,  )

										((_lowerCAmelCase)			,		)							=							result_with_labels.to_tuple()

										_lowerCAmelCase							=							model(lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__						)

										((_lowerCAmelCase)			,		)							=							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   SCREAMING_SNAKE_CASE__     (       self	:	Tuple  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	str  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	List[str]  ,  lowercase__	:	str  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Optional[Any]  ,  ):
										_lowerCAmelCase							=							FlaubertForSequenceClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__						)
										_lowerCAmelCase							=							model(lowercase__  ,  labels=lowercase__						)

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict  ,  lowercase__	:	List[str]  ,  lowercase__	:	int  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	Tuple  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Tuple  ,  lowercase__	:	Tuple  ,  lowercase__	:	List[str]  ,  ):
										_lowerCAmelCase							=							self.num_labels
										_lowerCAmelCase							=							FlaubertForTokenClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.num_labels)						)





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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_lowerCAmelCase							=							self.prepare_config_and_inputs()
										(
										    (
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		(
										    _lowerCAmelCase
										)			,		
										)							=							config_and_inputs
										_lowerCAmelCase							=							{
										    'input_ids': input_ids,
										    'token_type_ids': token_type_ids,
										    'lengths': input_lengths,
										    'attention_mask': input_mask,
										}
										return config, inputs_dict







@require_torch
class        lowerCamelCase__				(   UpperCAmelCase       ,UpperCAmelCase       ,unittest.TestCase				):
			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   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	int  ,  lowercase__	:	str  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int						):
										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   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	str  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Dict=False						):
										_lowerCAmelCase							=							super()._prepare_for_class(lowercase__  ,  lowercase__  ,  return_labels=lowercase__						)

										if return_labels:
																	if model_class.__name__ == "FlaubertForQuestionAnswering":
																								_lowerCAmelCase							=							torch.zeros(
																								    self.model_tester.batch_size  ,  dtype=torch.long  ,  device=lowercase__						)
																								_lowerCAmelCase							=							torch.zeros(
																								    self.model_tester.batch_size  ,  dtype=torch.long  ,  device=lowercase__						)

										return inputs_dict





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							FlaubertModelTester(self						)
										_lowerCAmelCase							=							ConfigTester(self  ,  config_class=lowercase__  ,  emb_dim=37						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Union[str, Any]						):
										self.config_tester.run_common_tests()





			def   SCREAMING_SNAKE_CASE__     (       self	:	Tuple						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_flaubert_model(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_flaubert_lm_head(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Union[str, Any]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_flaubert_simple_qa(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_flaubert_qa(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_flaubert_sequence_classif(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_flaubert_token_classif(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_flaubert_multiple_choice(*lowercase__						)





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	Union[str, Any]						):
										for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_lowerCAmelCase							=							FlaubertModel.from_pretrained(lowercase__						)
																	self.assertIsNotNone(lowercase__						)





			@slow
			@require_torch_gpu
			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										_lowerCAmelCase			,		_lowerCAmelCase							=							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

																	_lowerCAmelCase							=							True
																	_lowerCAmelCase							=							model_class(config=lowercase__						)

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

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







@require_torch
class        lowerCamelCase__				(   unittest.TestCase				):





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										_lowerCAmelCase							=							FlaubertModel.from_pretrained('flaubert/flaubert_base_cased'						)
										_lowerCAmelCase							=							torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]						)
										with torch.no_grad():
																	_lowerCAmelCase							=							model(lowercase__						)[0]
										_lowerCAmelCase							=							torch.Size((1, 11, 7_68)						)
										self.assertEqual(output.shape  ,  lowercase__						)
										_lowerCAmelCase							=							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]  ,  lowercase__  ,  atol=1e-4						)						)







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


_lowercase:	Optional[Any]       			=		{
    '''configuration_rag''': ['''RagConfig'''],
    '''retrieval_rag''': ['''RagRetriever'''],
    '''tokenization_rag''': ['''RagTokenizer'''],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowercase:	Optional[Any]       			=		[
				    '''RagModel''',
				    '''RagPreTrainedModel''',
				    '''RagSequenceForGeneration''',
				    '''RagTokenForGeneration''',
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowercase:	Optional[Any]       			=		[
				    '''TFRagModel''',
				    '''TFRagPreTrainedModel''',
				    '''TFRagSequenceForGeneration''',
				    '''TFRagTokenForGeneration''',
				]


if TYPE_CHECKING:
				from .configuration_rag import RagConfig
				from .retrieval_rag import RagRetriever
				from .tokenization_rag import RagTokenizer

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_rag import (
								    TFRagModel,
								    TFRagPreTrainedModel,
								    TFRagSequenceForGeneration,
								    TFRagTokenForGeneration,
								)

else:
				import sys

				_lowercase:	Optional[int]       			=		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






