





"""simple docstring"""





import argparse
import re
from pathlib import Path

import requests
import torch
from PIL import Image
from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor

from transformers import (
    EfficientFormerConfig,
    EfficientFormerForImageClassificationWithTeacher,
    EfficientFormerImageProcessor,
)
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling


def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						Tuple							,						SCREAMING_SNAKE_CASE     :						Tuple			):







      '''simple docstring'''




      lowerCAmelCase  						=       old_name

      if "patch_embed" in old_name:
            lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  						=       old_name.split("""."""			)

            if layer == "0":
                  lowerCAmelCase  						=       old_name.replace("""0"""							,						"""convolution1"""			)
            elif layer == "1":
                  lowerCAmelCase  						=       old_name.replace("""1"""							,						"""batchnorm_before"""			)
            elif layer == "3":
                  lowerCAmelCase  						=       old_name.replace("""3"""							,						"""convolution2"""			)
            else:
                  lowerCAmelCase  						=       old_name.replace("""4"""							,						"""batchnorm_after"""			)

      if "network" in old_name and re.search(R"""\d\.\d"""							,						SCREAMING_SNAKE_CASE			):
            lowerCAmelCase  						=       R"""\b\d{2}\b"""
            if bool(re.search(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE			)			):
                  lowerCAmelCase  						=       re.search(R"""\d\.\d\d."""							,						SCREAMING_SNAKE_CASE			).group()
            else:
                  lowerCAmelCase  						=       re.search(R"""\d\.\d."""							,						SCREAMING_SNAKE_CASE			).group()
            if int(match[0]			) < 6:
                  lowerCAmelCase  						=       old_name.replace(SCREAMING_SNAKE_CASE							,						""""""			)
                  lowerCAmelCase  						=       trimmed_name.replace("""network"""							,						match[0] + """.meta4D_layers.blocks.""" + match[2:-1]			)
                  lowerCAmelCase  						=       """intermediate_stages.""" + trimmed_name
            else:
                  lowerCAmelCase  						=       old_name.replace(SCREAMING_SNAKE_CASE							,						""""""			)
                  if int(match[2]			) < num_meta4D_last_stage:
                        lowerCAmelCase  						=       trimmed_name.replace("""network"""							,						"""meta4D_layers.blocks.""" + match[2]			)
                  else:
                        lowerCAmelCase  						=       str(int(match[2]			) - num_meta4D_last_stage			)
                        lowerCAmelCase  						=       trimmed_name.replace("""network"""							,						"""meta3D_layers.blocks.""" + layer_index			)
                        if "norm1" in old_name:
                              lowerCAmelCase  						=       trimmed_name.replace("""norm1"""							,						"""layernorm1"""			)
                        elif "norm2" in old_name:
                              lowerCAmelCase  						=       trimmed_name.replace("""norm2"""							,						"""layernorm2"""			)
                        elif "fc1" in old_name:
                              lowerCAmelCase  						=       trimmed_name.replace("""fc1"""							,						"""linear_in"""			)
                        elif "fc2" in old_name:
                              lowerCAmelCase  						=       trimmed_name.replace("""fc2"""							,						"""linear_out"""			)

                  lowerCAmelCase  						=       """last_stage.""" + trimmed_name

      elif "network" in old_name and re.search(R""".\d."""							,						SCREAMING_SNAKE_CASE			):
            lowerCAmelCase  						=       old_name.replace("""network"""							,						"""intermediate_stages"""			)

      if "fc" in new_name:
            lowerCAmelCase  						=       new_name.replace("""fc"""							,						"""convolution"""			)
      elif ("norm1" in new_name) and ("layernorm1" not in new_name):
            lowerCAmelCase  						=       new_name.replace("""norm1"""							,						"""batchnorm_before"""			)
      elif ("norm2" in new_name) and ("layernorm2" not in new_name):
            lowerCAmelCase  						=       new_name.replace("""norm2"""							,						"""batchnorm_after"""			)
      if "proj" in new_name:
            lowerCAmelCase  						=       new_name.replace("""proj"""							,						"""projection"""			)
      if "dist_head" in new_name:
            lowerCAmelCase  						=       new_name.replace("""dist_head"""							,						"""distillation_classifier"""			)
      elif "head" in new_name:
            lowerCAmelCase  						=       new_name.replace("""head"""							,						"""classifier"""			)
      elif "patch_embed" in new_name:
            lowerCAmelCase  						=       """efficientformer.""" + new_name
      elif new_name == "norm.weight" or new_name == "norm.bias":
            lowerCAmelCase  						=       new_name.replace("""norm"""							,						"""layernorm"""			)
            lowerCAmelCase  						=       """efficientformer.""" + new_name
      else:
            lowerCAmelCase  						=       """efficientformer.encoder.""" + new_name

      return new_name


def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						Dict							,						SCREAMING_SNAKE_CASE     :						Union[str, Any]			):







      '''simple docstring'''




      for key in checkpoint.copy().keys():
            lowerCAmelCase  						=       checkpoint.pop(SCREAMING_SNAKE_CASE			)
            lowerCAmelCase  						=       val

      return checkpoint


def 			UpperCAmelCase__							(  ):







      '''simple docstring'''




      lowerCAmelCase  						=       """http://images.cocodataset.org/val2017/000000039769.jpg"""
      lowerCAmelCase  						=       Image.open(requests.get(SCREAMING_SNAKE_CASE							,						stream=SCREAMING_SNAKE_CASE			).raw			)

      return image


def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						Path							,						SCREAMING_SNAKE_CASE     :						Path							,						SCREAMING_SNAKE_CASE     :						Path							,						SCREAMING_SNAKE_CASE     :						bool			):







      '''simple docstring'''




      lowerCAmelCase  						=       torch.load(SCREAMING_SNAKE_CASE							,						map_location="""cpu"""			)["""model"""]
      lowerCAmelCase  						=       EfficientFormerConfig.from_json_file(SCREAMING_SNAKE_CASE			)
      lowerCAmelCase  						=       EfficientFormerForImageClassificationWithTeacher(SCREAMING_SNAKE_CASE			)
      lowerCAmelCase  						=       """_""".join(checkpoint_path.split("""/"""			)[-1].split("""."""			)[0].split("""_"""			)[:-1]			)

      lowerCAmelCase  						=       config.depths[-1] - config.num_metaad_blocks + 1
      lowerCAmelCase  						=       convert_torch_checkpoint(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE			)

      model.load_state_dict(SCREAMING_SNAKE_CASE			)
      model.eval()

      lowerCAmelCase  						=       {
          """bilinear""": PILImageResampling.BILINEAR,
          """bicubic""": PILImageResampling.BICUBIC,
          """nearest""": PILImageResampling.NEAREST,
      }

      # prepare image
      lowerCAmelCase  						=       prepare_img()
      lowerCAmelCase  						=       2_56
      lowerCAmelCase  						=       2_24
      lowerCAmelCase  						=       EfficientFormerImageProcessor(
          size={"""shortest_edge""": image_size}							,						crop_size={"""height""": crop_size, """width""": crop_size}							,						resample=pillow_resamplings["""bicubic"""]							,						)
      lowerCAmelCase  						=       processor(images=SCREAMING_SNAKE_CASE							,						return_tensors="""pt"""			).pixel_values

      # original processing pipeline
      lowerCAmelCase  						=       Compose(
          [
              Resize(SCREAMING_SNAKE_CASE							,						interpolation=pillow_resamplings["""bicubic"""]			),
              CenterCrop(SCREAMING_SNAKE_CASE			),
              ToTensor(),
              Normalize(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE			),
          ]			)
      lowerCAmelCase  						=       image_transforms(SCREAMING_SNAKE_CASE			).unsqueeze(0			)

      assert torch.allclose(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE			)

      lowerCAmelCase  						=       model(SCREAMING_SNAKE_CASE			)
      lowerCAmelCase  						=       outputs.logits

      lowerCAmelCase  						=       (1, 10_00)

      if "l1" in model_name:
            lowerCAmelCase  						=       torch.Tensor(
                [-0.13_12, 0.43_53, -1.04_99, -0.51_24, 0.41_83, -0.67_93, -1.37_77, -0.08_93, -0.73_58, -2.43_28]			)
            assert torch.allclose(logits[0, :10]							,						SCREAMING_SNAKE_CASE							,						atol=1e-3			)
            assert logits.shape == expected_shape
      elif "l3" in model_name:
            lowerCAmelCase  						=       torch.Tensor(
                [-1.31_50, -1.54_56, -1.25_56, -0.84_96, -0.71_27, -0.78_97, -0.97_28, -0.30_52, 0.37_51, -0.31_27]			)
            assert torch.allclose(logits[0, :10]							,						SCREAMING_SNAKE_CASE							,						atol=1e-3			)
            assert logits.shape == expected_shape
      elif "l7" in model_name:
            lowerCAmelCase  						=       torch.Tensor(
                [-1.02_83, -1.41_31, -0.56_44, -1.31_15, -0.57_85, -1.20_49, -0.75_28, 0.19_92, -0.38_22, -0.08_78]			)
            assert logits.shape == expected_shape
      else:
            raise ValueError(
                F'Unknown model checkpoint: {checkpoint_path}. Supported version of efficientformer are l1, l3 and l7'			)

      # Save Checkpoints
      Path(SCREAMING_SNAKE_CASE			).mkdir(exist_ok=SCREAMING_SNAKE_CASE			)
      model.save_pretrained(SCREAMING_SNAKE_CASE			)
      print(F'Checkpoint successfuly converted. Model saved at {pytorch_dump_path}'			)
      processor.save_pretrained(SCREAMING_SNAKE_CASE			)
      print(F'Processor successfuly saved at {pytorch_dump_path}'			)

      if push_to_hub:
            print("""Pushing model to the hub..."""			)

            model.push_to_hub(
                repo_id=F'Bearnardd/{pytorch_dump_path}'							,						commit_message="""Add model"""							,						use_temp_dir=SCREAMING_SNAKE_CASE							,						)
            processor.push_to_hub(
                repo_id=F'Bearnardd/{pytorch_dump_path}'							,						commit_message="""Add image processor"""							,						use_temp_dir=SCREAMING_SNAKE_CASE							,						)


if __name__ == "__main__":
       SCREAMING_SNAKE_CASE__					=   argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           "--pytorch_model_path",
           default=None,
           type=str,
           required=True,
           help="Path to EfficientFormer pytorch checkpoint.",
       )
       parser.add_argument(
           "--config_file",
           default=None,
           type=str,
           required=True,
           help="The json file for EfficientFormer model config.",
       )
       parser.add_argument(
           "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
       )

       parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")
       parser.add_argument(
           "--no-push_to_hub",
           dest="push_to_hub",
           action="store_false",
           help="Do not push model and image processor to the hub",
       )
       parser.set_defaults(push_to_hub=True)

       SCREAMING_SNAKE_CASE__					=   parser.parse_args()
       convert_efficientformer_checkpoint(
           checkpoint_path=args.pytorch_model_path,
           efficientformer_config_file=args.config_file,
           pytorch_dump_path=args.pytorch_dump_path,
           push_to_hub=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 ) )






