




"""simple docstring"""




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     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a=13		, _a=7		, _a=True		, _a=True		, _a=True		, _a=True		, _a=True		, _a=False		, _a=False		, _a=False		, _a=2		, _a=99		, _a=0		, _a=32		, _a=5		, _a=4		, _a=0.1		, _a=0.1		, _a=512		, _a=12		, _a=2		, _a=0.02		, _a=3		, _a=4		, _a="last"		, _a=None		, _a=None		, ):





							"""simple docstring"""
							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    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							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    _lowerCAmelCase				(       self ):





							"""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    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaubertModel(config=_a )
							model.to(_a )
							model.eval()
							lowerCamelCase			      =					model(_a		, lengths=_a		, langs=_a )
							lowerCamelCase			      =					model(_a		, langs=_a )
							lowerCamelCase			      =					model(_a )
							self.parent.assertEqual(result.last_hidden_state.shape		, (self.batch_size, self.seq_length, self.hidden_size) )





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaubertWithLMHeadModel(_a )
							model.to(_a )
							model.eval()

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





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaubertForQuestionAnsweringSimple(_a )
							model.to(_a )
							model.eval()

							lowerCamelCase			      =					model(_a )

							lowerCamelCase			      =					model(_a		, start_positions=_a		, end_positions=_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    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaubertForQuestionAnswering(_a )
							model.to(_a )
							model.eval()

							lowerCamelCase			      =					model(_a )

							lowerCamelCase			      =					model(
							    _a		, start_positions=_a		, end_positions=_a		, cls_index=_a		, is_impossible=_a		, p_mask=_a		, )

							lowerCamelCase			      =					model(
							    _a		, start_positions=_a		, end_positions=_a		, cls_index=_a		, is_impossible=_a		, )

							((lowerCamelCase)      ,   )			      =					result_with_labels.to_tuple()

							lowerCamelCase			      =					model(_a		, start_positions=_a		, end_positions=_a )

							((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    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaubertForSequenceClassification(_a )
							model.to(_a )
							model.eval()

							lowerCamelCase			      =					model(_a )
							lowerCamelCase			      =					model(_a		, labels=_a )

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





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					self.num_labels
							lowerCamelCase			      =					FlaubertForTokenClassification(_a )
							model.to(_a )
							model.eval()

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





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, _a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					self.num_choices
							lowerCamelCase			      =					FlaubertForMultipleChoice(config=_a )
							model.to(_a )
							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(
							    _a		, attention_mask=_a		, token_type_ids=_a		, labels=_a		, )
							self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_choices) )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							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     __magic_name__							(       UpperCAmelCase__						,					UpperCAmelCase__						,					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    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a ):





							"""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    _lowerCAmelCase				(       self		, _a		, _a		, _a=False ):





							"""simple docstring"""
							lowerCamelCase			      =					super()._prepare_for_class(_a		, _a		, return_labels=_a )

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

							return inputs_dict





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaubertModelTester(self )
							lowerCamelCase			      =					ConfigTester(self		, config_class=_a		, emb_dim=37 )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							self.config_tester.run_common_tests()





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_model(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_lm_head(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_simple_qa(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_qa(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_sequence_classif(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_token_classif(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_flaubert_multiple_choice(*_a )





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCamelCase			      =					FlaubertModel.from_pretrained(_a )
										self.assertIsNotNone(_a )





				@slow
				@require_torch_gpu
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							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=_a )

										lowerCamelCase			      =					self._prepare_for_class(_a		, _a )
										lowerCamelCase			      =					torch.jit.trace(
										    _a		, (inputs_dict["""input_ids"""].to("""cpu""" ), inputs_dict["""attention_mask"""].to("""cpu""" )) )

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






@require_torch
class     __magic_name__							(       unittest.TestCase				):






				'''simple docstring'''





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaubertModel.from_pretrained("""flaubert/flaubert_base_cased""" )
							lowerCamelCase			      =					torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]] )
							with torch.no_grad():
										lowerCamelCase			      =					model(_a )[0]
							lowerCamelCase			      =					torch.Size((1, 11, 768) )
							self.assertEqual(output.shape		, _a )
							lowerCamelCase			      =					torch.tensor(
							    [[[-2.6_251, -1.4_298, -0.0_227], [-2.8_510, -1.6_387, 0.2_258], [-2.8_114, -1.1_832, -0.3_066]]] )

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









"""simple docstring"""




from math import ceil



def     a__ (				snake_case__       , snake_case__				)						->					Optional[int]:
			lowerCamelCase			      =					list(range(0       , snake_case__				)				)

			lowerCamelCase			      =					[item for sublist in list(device_map.values()				) for item in sublist]

			# Duplicate check
			lowerCamelCase			      =					[]
			for i in device_map_blocks:
						if device_map_blocks.count(snake_case__				) > 1 and i not in duplicate_blocks:
									duplicate_blocks.append(snake_case__				)
    # Missing blocks
			lowerCamelCase			      =					[i for i in blocks if i not in device_map_blocks]
			lowerCamelCase			      =					[i for i in device_map_blocks if i not in blocks]

			if len(snake_case__				) != 0:
						raise ValueError(
						    """Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device."""
						    """ These attention blocks were specified more than once: """ + str(snake_case__				)				)
			if len(snake_case__				) != 0:
						raise ValueError(
						    """There are attention blocks for this model that are not specified in the device_map. Add these attention """
						    """blocks to a device on the device_map: """ + str(snake_case__				)				)
			if len(snake_case__				) != 0:
						raise ValueError(
						    """The device_map contains more attention blocks than this model has. Remove these from the device_map:"""
						    + str(snake_case__				)				)



def     a__ (				snake_case__       , snake_case__				)						->					List[Any]:
			lowerCamelCase			      =					list(range(snake_case__				)				)
			lowerCamelCase			      =					int(ceil(n_layers / len(snake_case__				)				)				)
			lowerCamelCase			      =					[layers[i : i + n_blocks] for i in range(0       , snake_case__       , snake_case__				)]

			return dict(zip(snake_case__       , snake_case__				)				)



