






"""simple docstring"""

import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
lowercase__          = logging.get_logger(__name__)

def 		__lowerCamelCase	(  __UpperCamelCase       )						->    List[str]:

		"""simple docstring"""







		lowerCAmelCase_   :							List[Any]									=			YolosConfig()

		# size of the architecture
		if "yolos_ti" in yolos_name:
				lowerCAmelCase_   :							str									=			192
				lowerCAmelCase_   :							Tuple									=			768
				lowerCAmelCase_   :							Union[str, Any]									=			12
				lowerCAmelCase_   :							str									=			3
				lowerCAmelCase_   :							Union[str, Any]									=			[800, 1333]
				lowerCAmelCase_   :							int									=			False
		elif yolos_name == "yolos_s_dWr":
				lowerCAmelCase_   :							List[Any]									=			330
				lowerCAmelCase_   :							int									=			14
				lowerCAmelCase_   :							Tuple									=			6
				lowerCAmelCase_   :							Tuple									=			1320
		elif "yolos_s" in yolos_name:
				lowerCAmelCase_   :							Dict									=			384
				lowerCAmelCase_   :							Optional[int]									=			1536
				lowerCAmelCase_   :							Optional[int]									=			12
				lowerCAmelCase_   :							Any									=			6
		elif "yolos_b" in yolos_name:
				lowerCAmelCase_   :							Any									=			[800, 1344]

		lowerCAmelCase_   :							Any									=			91
		lowerCAmelCase_   :							Optional[Any]									=			"huggingface/label-files"
		lowerCAmelCase_   :							List[Any]									=			"coco-detection-id2label.json"
		lowerCAmelCase_   :							Tuple									=			json.load(open(hf_hub_download(__UpperCamelCase      ,       __UpperCamelCase      ,       repo_type="dataset"       )      ,       "r"       )       )
		lowerCAmelCase_   :							Union[str, Any]									=			{int(__UpperCamelCase       ): v for k, v in idalabel.items()}
		lowerCAmelCase_   :							int									=			idalabel
		lowerCAmelCase_   :							List[Any]									=			{v: k for k, v in idalabel.items()}

		return config

def 		__lowerCamelCase	(  __UpperCamelCase      ,       __UpperCamelCase      ,       __UpperCamelCase = False       )						->    Union[str, Any]:

		"""simple docstring"""







		for i in range(config.num_hidden_layers       ):
				# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
				lowerCAmelCase_   :							List[str]									=			state_dict.pop(f'''blocks.{i}.attn.qkv.weight'''       )
				lowerCAmelCase_   :							List[str]									=			state_dict.pop(f'''blocks.{i}.attn.qkv.bias'''       )
				# next, add query, keys and values (in that order) to the state dict
				lowerCAmelCase_   :							int									=			in_proj_weight[: config.hidden_size, :]
				lowerCAmelCase_   :							Any									=			in_proj_bias[: config.hidden_size]
				lowerCAmelCase_   :							List[str]									=			in_proj_weight[
				    config.hidden_size : config.hidden_size * 2, :
				]
				lowerCAmelCase_   :							Tuple									=			in_proj_bias[
				    config.hidden_size : config.hidden_size * 2
				]
				lowerCAmelCase_   :							Tuple									=			in_proj_weight[-config.hidden_size :, :]
				lowerCAmelCase_   :							Dict									=			in_proj_bias[-config.hidden_size :]

def 		__lowerCamelCase	(  __UpperCamelCase       )						->    str:

		"""simple docstring"""







		if "backbone" in name:
				lowerCAmelCase_   :							List[Any]									=			name.replace("backbone"      ,       "vit"       )
		if "cls_token" in name:
				lowerCAmelCase_   :							Union[str, Any]									=			name.replace("cls_token"      ,       "embeddings.cls_token"       )
		if "det_token" in name:
				lowerCAmelCase_   :							str									=			name.replace("det_token"      ,       "embeddings.detection_tokens"       )
		if "mid_pos_embed" in name:
				lowerCAmelCase_   :							Any									=			name.replace("mid_pos_embed"      ,       "encoder.mid_position_embeddings"       )
		if "pos_embed" in name:
				lowerCAmelCase_   :							Tuple									=			name.replace("pos_embed"      ,       "embeddings.position_embeddings"       )
		if "patch_embed.proj" in name:
				lowerCAmelCase_   :							List[Any]									=			name.replace("patch_embed.proj"      ,       "embeddings.patch_embeddings.projection"       )
		if "blocks" in name:
				lowerCAmelCase_   :							Optional[int]									=			name.replace("blocks"      ,       "encoder.layer"       )
		if "attn.proj" in name:
				lowerCAmelCase_   :							str									=			name.replace("attn.proj"      ,       "attention.output.dense"       )
		if "attn" in name:
				lowerCAmelCase_   :							Dict									=			name.replace("attn"      ,       "attention.self"       )
		if "norm1" in name:
				lowerCAmelCase_   :							Dict									=			name.replace("norm1"      ,       "layernorm_before"       )
		if "norm2" in name:
				lowerCAmelCase_   :							List[str]									=			name.replace("norm2"      ,       "layernorm_after"       )
		if "mlp.fc1" in name:
				lowerCAmelCase_   :							Tuple									=			name.replace("mlp.fc1"      ,       "intermediate.dense"       )
		if "mlp.fc2" in name:
				lowerCAmelCase_   :							Optional[int]									=			name.replace("mlp.fc2"      ,       "output.dense"       )
		if "class_embed" in name:
				lowerCAmelCase_   :							Union[str, Any]									=			name.replace("class_embed"      ,       "class_labels_classifier"       )
		if "bbox_embed" in name:
				lowerCAmelCase_   :							Dict									=			name.replace("bbox_embed"      ,       "bbox_predictor"       )
		if "vit.norm" in name:
				lowerCAmelCase_   :							int									=			name.replace("vit.norm"      ,       "vit.layernorm"       )

		return name

def 		__lowerCamelCase	(  __UpperCamelCase      ,       __UpperCamelCase       )						->    int:

		"""simple docstring"""







		for key in orig_state_dict.copy().keys():
				lowerCAmelCase_   :							Tuple									=			orig_state_dict.pop(__UpperCamelCase       )

				if "qkv" in key:
						lowerCAmelCase_   :							Optional[int]									=			key.split("."       )
						lowerCAmelCase_   :							Union[str, Any]									=			int(key_split[2]       )
						lowerCAmelCase_   :							List[str]									=			model.vit.encoder.layer[layer_num].attention.attention.all_head_size
						if "weight" in key:
								lowerCAmelCase_   :							int									=			val[:dim, :]
								lowerCAmelCase_   :							int									=			val[
								    dim : dim * 2, :
								]
								lowerCAmelCase_   :							Optional[int]									=			val[-dim:, :]
						else:
								lowerCAmelCase_   :							Tuple									=			val[:dim]
								lowerCAmelCase_   :							Optional[Any]									=			val[dim : dim * 2]
								lowerCAmelCase_   :							Optional[Any]									=			val[-dim:]
				else:
						lowerCAmelCase_   :							str									=			val

		return orig_state_dict

def 		__lowerCamelCase	(  )						->    List[str]:

		"""simple docstring"""







		lowerCAmelCase_   :							Any									=			"http://images.cocodataset.org/val2017/000000039769.jpg"
		lowerCAmelCase_   :							Optional[Any]									=			Image.open(requests.get(__UpperCamelCase      ,       stream=__UpperCamelCase       ).raw       )
		return im

@torch.no_grad()
def 		__lowerCamelCase	(  __UpperCamelCase      ,       __UpperCamelCase      ,       __UpperCamelCase      ,       __UpperCamelCase = False       )						->    str:

		"""simple docstring"""







		lowerCAmelCase_   :							Any									=			get_yolos_config(__UpperCamelCase       )

		# load original state_dict
		lowerCAmelCase_   :							Any									=			torch.load(__UpperCamelCase      ,       map_location="cpu"       )["model"]

		# load 🤗 model
		lowerCAmelCase_   :							Dict									=			YolosForObjectDetection(__UpperCamelCase       )
		model.eval()
		lowerCAmelCase_   :							int									=			convert_state_dict(__UpperCamelCase      ,       __UpperCamelCase       )
		model.load_state_dict(__UpperCamelCase       )

		# Check outputs on an image, prepared by YolosImageProcessor
		lowerCAmelCase_   :							Optional[int]									=			800 if yolos_name != "yolos_ti" else 512
		lowerCAmelCase_   :							Tuple									=			YolosImageProcessor(format="coco_detection"      ,       size=__UpperCamelCase       )
		lowerCAmelCase_   :							str									=			image_processor(images=prepare_img()      ,       return_tensors="pt"       )
		lowerCAmelCase_   :							Optional[int]									=			model(**__UpperCamelCase       )
		lowerCAmelCase_				,	lowerCAmelCase_   :							Optional[int]									=			outputs.logits, outputs.pred_boxes

		lowerCAmelCase_				,	lowerCAmelCase_   :							List[Any]									=			None, None
		if yolos_name == "yolos_ti":
				lowerCAmelCase_   :							List[str]									=			torch.tensor(
				    [[-39.50_22, -11.98_20, -17.68_88], [-29.95_74, -9.97_69, -17.76_91], [-42.32_81, -20.72_00, -30.62_94]]       )
				lowerCAmelCase_   :							List[Any]									=			torch.tensor(
				    [[0.40_21, 0.08_36, 0.79_79], [0.01_84, 0.26_09, 0.03_64], [0.17_81, 0.20_04, 0.20_95]]       )
		elif yolos_name == "yolos_s_200_pre":
				lowerCAmelCase_   :							Optional[int]									=			torch.tensor(
				    [[-24.02_48, -10.30_24, -14.82_90], [-42.03_92, -16.82_00, -27.43_34], [-27.27_43, -11.81_54, -18.71_48]]       )
				lowerCAmelCase_   :							Tuple									=			torch.tensor(
				    [[0.25_59, 0.54_55, 0.47_06], [0.29_89, 0.72_79, 0.18_75], [0.77_32, 0.40_17, 0.44_62]]       )
		elif yolos_name == "yolos_s_300_pre":
				lowerCAmelCase_   :							Any									=			torch.tensor(
				    [[-36.22_20, -14.43_85, -23.54_57], [-35.69_70, -14.75_83, -21.39_35], [-31.59_39, -13.60_42, -16.80_49]]       )
				lowerCAmelCase_   :							Tuple									=			torch.tensor(
				    [[0.76_14, 0.23_16, 0.47_28], [0.71_68, 0.44_95, 0.38_55], [0.49_96, 0.14_66, 0.99_96]]       )
		elif yolos_name == "yolos_s_dWr":
				lowerCAmelCase_   :							Union[str, Any]									=			torch.tensor(
				    [[-42.86_68, -24.10_49, -41.16_90], [-34.74_56, -14.12_74, -24.91_94], [-33.78_98, -12.19_46, -25.64_95]]       )
				lowerCAmelCase_   :							Optional[Any]									=			torch.tensor(
				    [[0.55_87, 0.27_73, 0.06_05], [0.50_04, 0.30_14, 0.99_94], [0.49_99, 0.15_48, 0.99_94]]       )
		elif yolos_name == "yolos_base":
				lowerCAmelCase_   :							Union[str, Any]									=			torch.tensor(
				    [[-40.60_64, -24.30_84, -32.64_47], [-55.19_90, -30.77_19, -35.58_77], [-51.43_11, -33.35_07, -35.64_62]]       )
				lowerCAmelCase_   :							int									=			torch.tensor(
				    [[0.55_55, 0.27_94, 0.06_55], [0.90_49, 0.26_64, 0.18_94], [0.91_83, 0.19_84, 0.16_35]]       )
		else:
				raise ValueError(f'''Unknown yolos_name: {yolos_name}'''       )

		assert torch.allclose(logits[0, :3, :3]      ,       __UpperCamelCase      ,       atol=1e-4       )
		assert torch.allclose(pred_boxes[0, :3, :3]      ,       __UpperCamelCase      ,       atol=1e-4       )

		Path(__UpperCamelCase       ).mkdir(exist_ok=__UpperCamelCase       )
		print(f'''Saving model {yolos_name} to {pytorch_dump_folder_path}'''       )
		model.save_pretrained(__UpperCamelCase       )
		print(f'''Saving image processor to {pytorch_dump_folder_path}'''       )
		image_processor.save_pretrained(__UpperCamelCase       )

		if push_to_hub:
				lowerCAmelCase_   :							str									=			{
				    "yolos_ti": "yolos-tiny",
				    "yolos_s_200_pre": "yolos-small",
				    "yolos_s_300_pre": "yolos-small-300",
				    "yolos_s_dWr": "yolos-small-dwr",
				    "yolos_base": "yolos-base",
				}

				print("Pushing to the hub..."       )
				lowerCAmelCase_   :							Tuple									=			model_mapping[yolos_name]
				image_processor.push_to_hub(__UpperCamelCase      ,       organization="hustvl"       )
				model.push_to_hub(__UpperCamelCase      ,       organization="hustvl"       )


if __name__ == "__main__":
					lowercase__          = argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--yolos_name""",
					    default="""yolos_s_200_pre""",
					    type=str,
					    help=(
					        """Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"""
					        """ 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."""
					    ),
					)
					parser.add_argument(
					    """--checkpoint_path""", default=None, type=str, help="""Path to the original state dict (.pth file)."""
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
					)
					parser.add_argument(
					    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
					)

					lowercase__          = parser.parse_args()
					convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)







"""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 		lowercase					(					_UpperCAmelCase					):





      def __init__(       self      ,						lowercase      ,						lowercase=13      ,						lowercase=7      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=False      ,						lowercase=False      ,						lowercase=False      ,						lowercase=2      ,						lowercase=99      ,						lowercase=0      ,						lowercase=32      ,						lowercase=5      ,						lowercase=4      ,						lowercase=0.1      ,						lowercase=0.1      ,						lowercase=512      ,						lowercase=12      ,						lowercase=2      ,						lowercase=0.02      ,						lowercase=3      ,						lowercase=4      ,						lowercase="last"      ,						lowercase=None      ,						lowercase=None      ,						)		->     int:
            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     _snake_case			(       self )		->     int:
            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     _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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Any:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Tuple:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     str:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Dict:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     int:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     int:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Tuple:
            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     _snake_case			(       self )		->     List[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 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase     ,  unittest.TestCase					):
      _SCREAMING_SNAKE_CASE					=      (
          (
              FlaubertModel,
              FlaubertWithLMHeadModel,
              FlaubertForQuestionAnswering,
              FlaubertForQuestionAnsweringSimple,
              FlaubertForSequenceClassification,
              FlaubertForTokenClassification,
              FlaubertForMultipleChoice,
          )
          if is_torch_available()
          else ()
      )
      _SCREAMING_SNAKE_CASE					=      (
          {
              '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      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Optional[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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase=False )		->     Optional[Any]:
            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     _snake_case			(       self )		->     List[str]:
            lowerCAmelCase  						=       FlaubertModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase      ,						emb_dim=37 )





      def     _snake_case			(       self )		->     Tuple:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_model(*lowercase )





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_lm_head(*lowercase )





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_simple_qa(*lowercase )





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_qa(*lowercase )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_sequence_classif(*lowercase )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_token_classif(*lowercase )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_multiple_choice(*lowercase )





      @slow
      def     _snake_case			(       self )		->     Tuple:
            for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCAmelCase  						=       FlaubertModel.from_pretrained(lowercase )
                  self.assertIsNotNone(lowercase )





      @slow
      @require_torch_gpu
      def     _snake_case			(       self )		->     List[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 		lowercase					(					unittest.TestCase					):





      @slow
      def     _snake_case			(       self )		->     Optional[int]:
            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(lowercase )[0]
            lowerCAmelCase  						=       torch.Size((1, 11, 768) )
            self.assertEqual(output.shape      ,						lowercase )
            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]      ,						lowercase      ,						atol=1e-4 ) )






