





import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(42)

_lowerCamelCase 					=      'bert-base-cased'
_lowerCamelCase 					=      'fp16'
_lowerCamelCase 					=      'bf16'
_lowerCamelCase 					=      [FPaa, BFaa]




@require_fsdp
@require_cuda
class 	UpperCamelCase_							(  UpperCamelCase__     ):


							def   _snake_case	(				self     :Optional[Any]		) -> Optional[Any]:



									"""simple docstring"""



									super().setUp()

									SCREAMING_SNAKE_CASE__								=   dict(
									    ACCELERATE_USE_FSDP="""true"""				,					MASTER_ADDR="""localhost"""				,					MASTER_PORT="""10999"""				,					RANK="""0"""				,					LOCAL_RANK="""0"""				,					WORLD_SIZE="""1"""				,					)


							def   _snake_case	(				self     :List[Any]		) -> Tuple:



									"""simple docstring"""



									from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

									for i, strategy in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
											SCREAMING_SNAKE_CASE__								=   f'''{i + 1}'''
											SCREAMING_SNAKE_CASE__								=   strategy
											with mockenv_context(**__A		):
													SCREAMING_SNAKE_CASE__								=   FullyShardedDataParallelPlugin()
													self.assertEqual(fsdp_plugin.sharding_strategy				,					ShardingStrategy(i + 1		)		)


							def   _snake_case	(				self     :int		) -> List[str]:



									"""simple docstring"""



									from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

									for i, prefetch_policy in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
											SCREAMING_SNAKE_CASE__								=   prefetch_policy
											with mockenv_context(**__A		):
													SCREAMING_SNAKE_CASE__								=   FullyShardedDataParallelPlugin()
													if prefetch_policy == "NO_PREFETCH":
															self.assertIsNone(fsdp_plugin.backward_prefetch		)
													else:
															self.assertEqual(fsdp_plugin.backward_prefetch				,					BackwardPrefetch(i + 1		)		)


							def   _snake_case	(				self     :List[str]		) -> List[str]:



									"""simple docstring"""



									from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

									for i, state_dict_type in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
											SCREAMING_SNAKE_CASE__								=   state_dict_type
											with mockenv_context(**__A		):
													SCREAMING_SNAKE_CASE__								=   FullyShardedDataParallelPlugin()
													self.assertEqual(fsdp_plugin.state_dict_type				,					StateDictType(i + 1		)		)
													if state_dict_type == "FULL_STATE_DICT":
															self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu		)
															self.assertTrue(fsdp_plugin.state_dict_config.ranka_only		)


							def   _snake_case	(				self     :str		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   AutoModel.from_pretrained(__A		)
									for policy in FSDP_AUTO_WRAP_POLICY:
											SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
											SCREAMING_SNAKE_CASE__								=   policy
											if policy == "TRANSFORMER_BASED_WRAP":
													SCREAMING_SNAKE_CASE__								=   """BertLayer"""
											elif policy == "SIZE_BASED_WRAP":
													SCREAMING_SNAKE_CASE__								=   """2000"""
											with mockenv_context(**__A		):
													SCREAMING_SNAKE_CASE__								=   FullyShardedDataParallelPlugin()
													fsdp_plugin.set_auto_wrap_policy(__A		)
													if policy == "NO_WRAP":
															self.assertIsNone(fsdp_plugin.auto_wrap_policy		)
													else:
															self.assertIsNotNone(fsdp_plugin.auto_wrap_policy		)

									SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
									SCREAMING_SNAKE_CASE__								=   """TRANSFORMER_BASED_WRAP"""
									SCREAMING_SNAKE_CASE__								=   """T5Layer"""
									with mockenv_context(**__A		):
											SCREAMING_SNAKE_CASE__								=   FullyShardedDataParallelPlugin()
											with self.assertRaises(__A		) as cm:
													fsdp_plugin.set_auto_wrap_policy(__A		)
											self.assertTrue("""Could not find the transformer layer class to wrap in the model.""" in str(cm.exception		)		)

									SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
									SCREAMING_SNAKE_CASE__								=   """SIZE_BASED_WRAP"""
									SCREAMING_SNAKE_CASE__								=   """0"""
									with mockenv_context(**__A		):
											SCREAMING_SNAKE_CASE__								=   FullyShardedDataParallelPlugin()
											fsdp_plugin.set_auto_wrap_policy(__A		)
											self.assertIsNone(fsdp_plugin.auto_wrap_policy		)


							def   _snake_case	(				self     :Optional[Any]		) -> Optional[int]:



									"""simple docstring"""



									from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
									from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

									for mp_dtype in dtypes:
											SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
											SCREAMING_SNAKE_CASE__								=   mp_dtype
											with mockenv_context(**__A		):
													SCREAMING_SNAKE_CASE__								=   Accelerator()
													if mp_dtype == "fp16":
															SCREAMING_SNAKE_CASE__								=   torch.floataa
													elif mp_dtype == "bf16":
															SCREAMING_SNAKE_CASE__								=   torch.bfloataa
													SCREAMING_SNAKE_CASE__								=   MixedPrecision(param_dtype=__A				,					reduce_dtype=__A				,					buffer_dtype=__A		)
													self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy				,					__A		)
													if mp_dtype == FPaa:
															self.assertTrue(isinstance(accelerator.scaler				,					__A		)		)
													elif mp_dtype == BFaa:
															self.assertIsNone(accelerator.scaler		)
													AcceleratorState._reset_state(__A		)


							def   _snake_case	(				self     :str		) -> str:



									"""simple docstring"""



									from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

									for flag in [True, False]:
											SCREAMING_SNAKE_CASE__								=   self.dist_env.copy()
											SCREAMING_SNAKE_CASE__								=   str(__A		).lower()
											with mockenv_context(**__A		):
													SCREAMING_SNAKE_CASE__								=   FullyShardedDataParallelPlugin()
													self.assertEqual(fsdp_plugin.cpu_offload				,					CPUOffload(offload_params=__A		)		)




@require_fsdp
@require_multi_gpu
@slow
class 	UpperCamelCase_							(  UpperCamelCase__     ):


							def   _snake_case	(				self     :Any		) -> Any:



									"""simple docstring"""



									super().setUp()
									SCREAMING_SNAKE_CASE__								=   0.8_2
									SCREAMING_SNAKE_CASE__								=   [
									    """fsdp_shard_grad_op_transformer_based_wrap""",
									    """fsdp_full_shard_transformer_based_wrap""",
									]
									SCREAMING_SNAKE_CASE__								=   {
									    """multi_gpu_fp16""": 3200,
									    """fsdp_shard_grad_op_transformer_based_wrap_fp16""": 2000,
									    """fsdp_full_shard_transformer_based_wrap_fp16""": 1900,
									    # Disabling below test as it overwhelms the RAM memory usage
									    # on CI self-hosted runner leading to tests getting killed.
									    # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
									}
									SCREAMING_SNAKE_CASE__								=   160
									SCREAMING_SNAKE_CASE__								=   160

									SCREAMING_SNAKE_CASE__								=   inspect.getfile(accelerate.test_utils		)
									SCREAMING_SNAKE_CASE__								=   os.path.sep.join(mod_file.split(os.path.sep		)[:-1] + ["""scripts""", """external_deps"""]		)


							def   _snake_case	(				self     :Union[str, Any]		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   os.path.join(self.test_scripts_folder				,					"""test_performance.py"""		)
									SCREAMING_SNAKE_CASE__								=   ["""accelerate""", """launch""", """--num_processes=2""", """--num_machines=1""", """--machine_rank=0""", """--use_fsdp"""]
									for config in self.performance_configs:
											SCREAMING_SNAKE_CASE__								=   cmd.copy()
											for i, strategy in enumerate(__A		):
													if strategy.lower() in config:
															cmd_config.append(f'''--fsdp_sharding_strategy={i+1}'''		)
															break

											if "fp32" in config:
													cmd_config.append("""--mixed_precision=no"""		)
											else:
													cmd_config.append("""--mixed_precision=fp16"""		)

											if "cpu_offload" in config:
													cmd_config.append("""--fsdp_offload_params=True"""		)

											for policy in FSDP_AUTO_WRAP_POLICY:
													if policy.lower() in config:
															cmd_config.append(f'''--fsdp_auto_wrap_policy={policy}'''		)
															break

											if policy == "TRANSFORMER_BASED_WRAP":
													cmd_config.append("""--fsdp_transformer_layer_cls_to_wrap=BertLayer"""		)
											elif policy == "SIZE_BASED_WRAP":
													cmd_config.append("""--fsdp_min_num_params=2000"""		)

											cmd_config.extend(
											    [
											        self.test_file_path,
											        f'''--output_dir={self.tmpdir}''',
											        f'''--performance_lower_bound={self.performance_lower_bound}''',
											    ]		)
											with patch_environment(omp_num_threads=1		):
													execute_subprocess_async(__A				,					env=os.environ.copy()		)


							def   _snake_case	(				self     :Dict		) -> Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   os.path.join(self.test_scripts_folder				,					"""test_checkpointing.py"""		)
									SCREAMING_SNAKE_CASE__								=   [
									    """accelerate""",
									    """launch""",
									    """--num_processes=2""",
									    """--num_machines=1""",
									    """--machine_rank=0""",
									    """--use_fsdp""",
									    """--mixed_precision=fp16""",
									    """--fsdp_transformer_layer_cls_to_wrap=BertLayer""",
									]

									for i, strategy in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   cmd.copy()
											cmd_config.append(f'''--fsdp_sharding_strategy={i+1}'''		)
											if strategy != "FULL_SHARD":
													continue
											SCREAMING_SNAKE_CASE__								=   len(__A		)
											for state_dict_type in FSDP_STATE_DICT_TYPE:
													SCREAMING_SNAKE_CASE__								=   cmd_config[:state_dict_config_index]
													cmd_config.append(f'''--fsdp_state_dict_type={state_dict_type}'''		)
													cmd_config.extend(
													    [
													        self.test_file_path,
													        f'''--output_dir={self.tmpdir}''',
													        """--partial_train_epoch=1""",
													    ]		)
													with patch_environment(omp_num_threads=1		):
															execute_subprocess_async(__A				,					env=os.environ.copy()		)

													SCREAMING_SNAKE_CASE__								=   cmd_config[:-1]
													SCREAMING_SNAKE_CASE__								=   os.path.join(self.tmpdir				,					"""epoch_0"""		)
													cmd_config.extend(
													    [
													        f'''--resume_from_checkpoint={resume_from_checkpoint}''',
													    ]		)
													with patch_environment(omp_num_threads=1		):
															execute_subprocess_async(__A				,					env=os.environ.copy()		)


							def   _snake_case	(				self     :Tuple		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   os.path.join(self.test_scripts_folder				,					"""test_peak_memory_usage.py"""		)
									SCREAMING_SNAKE_CASE__								=   [
									    """accelerate""",
									    """launch""",
									    """--num_processes=2""",
									    """--num_machines=1""",
									    """--machine_rank=0""",
									]
									for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
											SCREAMING_SNAKE_CASE__								=   cmd.copy()
											if "fp16" in spec:
													cmd_config.extend(["""--mixed_precision=fp16"""]		)
											else:
													cmd_config.extend(["""--mixed_precision=no"""]		)

											if "multi_gpu" in spec:
													continue
											else:
													cmd_config.extend(["""--use_fsdp"""]		)
													for i, strategy in enumerate(__A		):
															if strategy.lower() in spec:
																	cmd_config.append(f'''--fsdp_sharding_strategy={i+1}'''		)
																	break

													if "cpu_offload" in spec:
															cmd_config.append("""--fsdp_offload_params=True"""		)

													for policy in FSDP_AUTO_WRAP_POLICY:
															if policy.lower() in spec:
																	cmd_config.append(f'''--fsdp_auto_wrap_policy={policy}'''		)
																	break

													if policy == "TRANSFORMER_BASED_WRAP":
															cmd_config.append("""--fsdp_transformer_layer_cls_to_wrap=BertLayer"""		)
													elif policy == "SIZE_BASED_WRAP":
															cmd_config.append("""--fsdp_min_num_params=2000"""		)

											cmd_config.extend(
											    [
											        self.test_file_path,
											        f'''--output_dir={self.tmpdir}''',
											        f'''--peak_memory_upper_bound={peak_mem_upper_bound}''',
											        f'''--n_train={self.n_train}''',
											        f'''--n_val={self.n_val}''',
											    ]		)
											with patch_environment(omp_num_threads=1		):
													execute_subprocess_async(__A				,					env=os.environ.copy()		)






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 	UpperCamelCase_							(  UpperCamelCase__     ):


							def __init__(				self     :Union[str, Any]				,					__A     :Optional[int]				,					__A     :Tuple=13				,					__A     :Dict=7				,					__A     :Dict=True				,					__A     :str=True				,					__A     :Optional[Any]=True				,					__A     :Optional[Any]=True				,					__A     :Optional[Any]=True				,					__A     :Any=False				,					__A     :Dict=False				,					__A     :Any=False				,					__A     :Tuple=2				,					__A     :Dict=99				,					__A     :Optional[Any]=0				,					__A     :List[str]=32				,					__A     :Optional[int]=5				,					__A     :Dict=4				,					__A     :List[str]=0.1				,					__A     :Union[str, Any]=0.1				,					__A     :Tuple=512				,					__A     :Any=12				,					__A     :Optional[int]=2				,					__A     :Union[str, Any]=0.0_2				,					__A     :Dict=3				,					__A     :Optional[int]=4				,					__A     :Any="last"				,					__A     :List[Any]=None				,					__A     :Any=None				,					) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   parent
									SCREAMING_SNAKE_CASE__								=   batch_size
									SCREAMING_SNAKE_CASE__								=   seq_length
									SCREAMING_SNAKE_CASE__								=   is_training
									SCREAMING_SNAKE_CASE__								=   use_input_lengths
									SCREAMING_SNAKE_CASE__								=   use_token_type_ids
									SCREAMING_SNAKE_CASE__								=   use_labels
									SCREAMING_SNAKE_CASE__								=   gelu_activation
									SCREAMING_SNAKE_CASE__								=   sinusoidal_embeddings
									SCREAMING_SNAKE_CASE__								=   causal
									SCREAMING_SNAKE_CASE__								=   asm
									SCREAMING_SNAKE_CASE__								=   n_langs
									SCREAMING_SNAKE_CASE__								=   vocab_size
									SCREAMING_SNAKE_CASE__								=   n_special
									SCREAMING_SNAKE_CASE__								=   hidden_size
									SCREAMING_SNAKE_CASE__								=   num_hidden_layers
									SCREAMING_SNAKE_CASE__								=   num_attention_heads
									SCREAMING_SNAKE_CASE__								=   hidden_dropout_prob
									SCREAMING_SNAKE_CASE__								=   attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE__								=   max_position_embeddings
									SCREAMING_SNAKE_CASE__								=   type_vocab_size
									SCREAMING_SNAKE_CASE__								=   type_sequence_label_size
									SCREAMING_SNAKE_CASE__								=   initializer_range
									SCREAMING_SNAKE_CASE__								=   num_labels
									SCREAMING_SNAKE_CASE__								=   num_choices
									SCREAMING_SNAKE_CASE__								=   summary_type
									SCREAMING_SNAKE_CASE__								=   use_proj
									SCREAMING_SNAKE_CASE__								=   scope


							def   _snake_case	(				self     :Optional[Any]		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   ids_tensor([self.batch_size, self.seq_length]				,					self.vocab_size		)
									SCREAMING_SNAKE_CASE__								=   random_attention_mask([self.batch_size, self.seq_length]		)

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

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

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

									SCREAMING_SNAKE_CASE__								=   self.get_config()

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


							def   _snake_case	(				self     :List[str]		) -> Optional[int]:



									"""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   _snake_case	(				self     :Tuple				,					__A     :str				,					__A     :int				,					__A     :Optional[int]				,					__A     :Any				,					__A     :Union[str, Any]				,					__A     :Optional[int]				,					__A     :Union[str, Any]				,					__A     :Union[str, Any]				,					__A     :str				,					) -> List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   FlaubertModel(config=__A		)
									model.to(__A		)
									model.eval()
									SCREAMING_SNAKE_CASE__								=   model(__A				,					lengths=__A				,					langs=__A		)
									SCREAMING_SNAKE_CASE__								=   model(__A				,					langs=__A		)
									SCREAMING_SNAKE_CASE__								=   model(__A		)
									self.parent.assertEqual(result.last_hidden_state.shape				,					(self.batch_size, self.seq_length, self.hidden_size)		)


							def   _snake_case	(				self     :str				,					__A     :Any				,					__A     :str				,					__A     :Union[str, Any]				,					__A     :Optional[Any]				,					__A     :Optional[int]				,					__A     :Any				,					__A     :Union[str, Any]				,					__A     :Optional[Any]				,					__A     :Union[str, Any]				,					) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   FlaubertWithLMHeadModel(__A		)
									model.to(__A		)
									model.eval()

									SCREAMING_SNAKE_CASE__								=   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   _snake_case	(				self     :Tuple				,					__A     :Union[str, Any]				,					__A     :Optional[Any]				,					__A     :Dict				,					__A     :Dict				,					__A     :Union[str, Any]				,					__A     :List[str]				,					__A     :Optional[int]				,					__A     :int				,					__A     :str				,					) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   FlaubertForQuestionAnsweringSimple(__A		)
									model.to(__A		)
									model.eval()

									SCREAMING_SNAKE_CASE__								=   model(__A		)

									SCREAMING_SNAKE_CASE__								=   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   _snake_case	(				self     :List[str]				,					__A     :Any				,					__A     :int				,					__A     :Tuple				,					__A     :Optional[Any]				,					__A     :Tuple				,					__A     :Optional[int]				,					__A     :str				,					__A     :int				,					__A     :str				,					) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   FlaubertForQuestionAnswering(__A		)
									model.to(__A		)
									model.eval()

									SCREAMING_SNAKE_CASE__								=   model(__A		)

									SCREAMING_SNAKE_CASE__								=   model(
									    __A				,					start_positions=__A				,					end_positions=__A				,					cls_index=__A				,					is_impossible=__A				,					p_mask=__A				,					)

									SCREAMING_SNAKE_CASE__								=   model(
									    __A				,					start_positions=__A				,					end_positions=__A				,					cls_index=__A				,					is_impossible=__A				,					)

									((SCREAMING_SNAKE_CASE__)						,    )								=   result_with_labels.to_tuple()

									SCREAMING_SNAKE_CASE__								=   model(__A				,					start_positions=__A				,					end_positions=__A		)

									((SCREAMING_SNAKE_CASE__)						,    )								=   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   _snake_case	(				self     :Optional[int]				,					__A     :str				,					__A     :Optional[int]				,					__A     :Tuple				,					__A     :Dict				,					__A     :List[str]				,					__A     :Tuple				,					__A     :List[str]				,					__A     :Dict				,					__A     :List[str]				,					) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   FlaubertForSequenceClassification(__A		)
									model.to(__A		)
									model.eval()

									SCREAMING_SNAKE_CASE__								=   model(__A		)
									SCREAMING_SNAKE_CASE__								=   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   _snake_case	(				self     :Optional[Any]				,					__A     :Optional[Any]				,					__A     :Optional[Any]				,					__A     :List[str]				,					__A     :Optional[Any]				,					__A     :int				,					__A     :Tuple				,					__A     :Optional[int]				,					__A     :Union[str, Any]				,					__A     :Dict				,					) -> Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.num_labels
									SCREAMING_SNAKE_CASE__								=   FlaubertForTokenClassification(__A		)
									model.to(__A		)
									model.eval()

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


							def   _snake_case	(				self     :str				,					__A     :Any				,					__A     :Tuple				,					__A     :List[str]				,					__A     :Tuple				,					__A     :Any				,					__A     :int				,					__A     :Dict				,					__A     :List[str]				,					__A     :Tuple				,					) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.num_choices
									SCREAMING_SNAKE_CASE__								=   FlaubertForMultipleChoice(config=__A		)
									model.to(__A		)
									model.eval()
									SCREAMING_SNAKE_CASE__								=   input_ids.unsqueeze(1		).expand(-1				,					self.num_choices				,					-1		).contiguous()
									SCREAMING_SNAKE_CASE__								=   token_type_ids.unsqueeze(1		).expand(-1				,					self.num_choices				,					-1		).contiguous()
									SCREAMING_SNAKE_CASE__								=   input_mask.unsqueeze(1		).expand(-1				,					self.num_choices				,					-1		).contiguous()
									SCREAMING_SNAKE_CASE__								=   model(
									    __A				,					attention_mask=__A				,					token_type_ids=__A				,					labels=__A				,					)
									self.parent.assertEqual(result.logits.shape				,					(self.batch_size, self.num_choices)		)


							def   _snake_case	(				self     :Union[str, Any]		) -> Any:



									"""simple docstring"""



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




@require_torch
class 	UpperCamelCase_							(  UpperCamelCase__					,	UpperCamelCase__					,	unittest.TestCase     ):
							lowerCamelCase_							     =       (
							    (
							        FlaubertModel,
							        FlaubertWithLMHeadModel,
							        FlaubertForQuestionAnswering,
							        FlaubertForQuestionAnsweringSimple,
							        FlaubertForSequenceClassification,
							        FlaubertForTokenClassification,
							        FlaubertForMultipleChoice,
							    )
							    if is_torch_available()
							    else ()
							)
							lowerCamelCase_							     =       (
							    {
							        "feature-extraction": FlaubertModel,
							        "fill-mask": FlaubertWithLMHeadModel,
							        "question-answering": FlaubertForQuestionAnsweringSimple,
							        "text-classification": FlaubertForSequenceClassification,
							        "token-classification": FlaubertForTokenClassification,
							        "zero-shot": FlaubertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)


							def   _snake_case	(				self     :Any				,					__A     :Optional[int]				,					__A     :Optional[int]				,					__A     :Dict				,					__A     :List[Any]				,					__A     :Tuple		) -> str:



									"""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     :Tuple				,					__A     :List[str]				,					__A     :Optional[int]				,					__A     :Dict=False		) -> List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   super()._prepare_for_class(__A				,					__A				,					return_labels=__A		)

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

									return inputs_dict


							def   _snake_case	(				self     :str		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   FlaubertModelTester(self		)
									SCREAMING_SNAKE_CASE__								=   ConfigTester(self				,					config_class=__A				,					emb_dim=37		)


							def   _snake_case	(				self     :int		) -> int:



									"""simple docstring"""



									self.config_tester.run_common_tests()


							def   _snake_case	(				self     :Optional[Any]		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_flaubert_model(*__A		)


							def   _snake_case	(				self     :Tuple		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_flaubert_lm_head(*__A		)


							def   _snake_case	(				self     :str		) -> Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_flaubert_simple_qa(*__A		)


							def   _snake_case	(				self     :Union[str, Any]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_flaubert_qa(*__A		)


							def   _snake_case	(				self     :str		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_flaubert_sequence_classif(*__A		)


							def   _snake_case	(				self     :Any		) -> Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_flaubert_token_classif(*__A		)


							def   _snake_case	(				self     :Any		) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_flaubert_multiple_choice(*__A		)


							@slow
							def   _snake_case	(				self     :Union[str, Any]		) -> List[str]:



									"""simple docstring"""



									for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											SCREAMING_SNAKE_CASE__								=   FlaubertModel.from_pretrained(__A		)
											self.assertIsNotNone(__A		)


							@slow
							@require_torch_gpu
							def   _snake_case	(				self     :Tuple		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   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

											SCREAMING_SNAKE_CASE__								=   True
											SCREAMING_SNAKE_CASE__								=   model_class(config=__A		)

											SCREAMING_SNAKE_CASE__								=   self._prepare_for_class(__A				,					__A		)
											SCREAMING_SNAKE_CASE__								=   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"""		)		)
													SCREAMING_SNAKE_CASE__								=   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 	UpperCamelCase_							(  unittest.TestCase     ):


							@slow
							def   _snake_case	(				self     :Dict		) -> Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   FlaubertModel.from_pretrained("""flaubert/flaubert_base_cased"""		)
									SCREAMING_SNAKE_CASE__								=   torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]		)
									with torch.no_grad():
											SCREAMING_SNAKE_CASE__								=   model(__A		)[0]
									SCREAMING_SNAKE_CASE__								=   torch.Size((1, 11, 768)		)
									self.assertEqual(output.shape				,					__A		)
									SCREAMING_SNAKE_CASE__								=   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]				,					__A				,					atol=1E-4		)		)