



import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import torch
from datasets import load_dataset
from torchvision.transforms import Compose, Lambda, Normalize, RandomHorizontalFlip, RandomResizedCrop, ToTensor

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    AutoModelForMaskedImageModeling,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
_lowercase						  =		logging.getLogger(__name__)

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.31.0")

require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/image-pretraining/requirements.txt")

_lowercase						  =		list(MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING.keys())
_lowercase						  =		tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)



@dataclass
class     _UpperCAmelCase			:
     UpperCamelCase__					    =				field(
         default='''cifar10'''			,      metadata={'''help''': '''Name of a dataset from the datasets package'''}     )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={'''help''': '''The configuration name of the dataset to use (via the datasets library).'''}     )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={'''help''': '''The column name of the images in the files. If not set, will try to use \'image\' or \'img\'.'''}			,      )
     UpperCamelCase__					    =				field(default=A__			,      metadata={'''help''': '''A folder containing the training data.'''}     )
     UpperCamelCase__					    =				field(default=A__			,      metadata={'''help''': '''A folder containing the validation data.'''}     )
     UpperCamelCase__					    =				field(
         default=0.15			,      metadata={'''help''': '''Percent to split off of train for validation.'''}     )
     UpperCamelCase__					    =				field(default=32			,      metadata={'''help''': '''The size of the square patches to use for masking.'''}     )
     UpperCamelCase__					    =				field(
         default=0.6			,      metadata={'''help''': '''Percentage of patches to mask.'''}			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={
             '''help''': (
                 '''For debugging purposes or quicker training, truncate the number of training examples to this '''
                 '''value if set.'''
             )
         }			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={
             '''help''': (
                 '''For debugging purposes or quicker training, truncate the number of evaluation examples to this '''
                 '''value if set.'''
             )
         }			,      )




     def 		snake_case_		(				self):
         A__		 =					{}
         if self.train_dir is not None:
             A__		 =					self.train_dir
         if self.validation_dir is not None:
             A__		 =					self.validation_dir
         A__		 =					data_files if data_files else None







@dataclass
class     _UpperCAmelCase			:
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={
             '''help''': (
                 '''The model checkpoint for weights initialization. Can be a local path to a pytorch_model.bin or a '''
                 '''checkpoint identifier on the hub. '''
                 '''Don\'t set if you want to train a model from scratch.'''
             )
         }			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={'''help''': '''If training from scratch, pass a model type from the list: ''' + ''', '''.join(A__     )}			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''}     )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={
             '''help''': (
                 '''Override some existing default config settings when a model is trained from scratch. Example: '''
                 '''n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index'''
             )
         }			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={'''help''': '''Where do you want to store (cache) the pretrained models/datasets downloaded from the hub'''}			,      )
     UpperCamelCase__					    =				field(
         default='''main'''			,      metadata={'''help''': '''The specific model version to use (can be a branch name, tag name or commit id).'''}			,      )
     UpperCamelCase__					    =				field(default=A__			,      metadata={'''help''': '''Name or path of preprocessor config.'''}     )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={
             '''help''': (
                 '''Will use the token generated when running `huggingface-cli login` (necessary to use this script '''
                 '''with private models).'''
             )
         }			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={
             '''help''': (
                 '''The size (resolution) of each image. If not specified, will use `image_size` of the configuration.'''
             )
         }			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={
             '''help''': (
                 '''The size (resolution) of each patch. If not specified, will use `patch_size` of the configuration.'''
             )
         }			,      )
     UpperCamelCase__					    =				field(
         default=A__			,      metadata={'''help''': '''Stride to use for the encoder.'''}			,      )







class     _UpperCAmelCase			:




     def __init__(				self ,						a__=1_9_2 ,						a__=3_2 ,						a__=4 ,						a__=0.6):
         A__		 =					input_size
         A__		 =					mask_patch_size
         A__		 =					model_patch_size
         A__		 =					mask_ratio

         if self.input_size % self.mask_patch_size != 0:
             raise ValueError('''Input size must be divisible by mask patch size''')
         if self.mask_patch_size % self.model_patch_size != 0:
             raise ValueError('''Mask patch size must be divisible by model patch size''')

         A__		 =					self.input_size // self.mask_patch_size
         A__		 =					self.mask_patch_size // self.model_patch_size

         A__		 =					self.rand_size**2
         A__		 =					int(np.ceil(self.token_count * self.mask_ratio))




     def __call__(				self):
         A__		 =					np.random.permutation(self.token_count)[: self.mask_count]
         A__		 =					np.zeros(self.token_count ,						dtype=a__)
         A__		 =					1

         A__		 =					mask.reshape((self.rand_size, self.rand_size))
         A__		 =					mask.repeat(self.scale ,						axis=0).repeat(self.scale ,						axis=1)

         return torch.tensor(mask.flatten())







def 		lowerCAmelCase__    (					UpperCamelCase_				:       Optional[Any]	)->  Any:
    A__		 =					torch.stack([example['''pixel_values'''] for example in examples]	)
    A__		 =					torch.stack([example['''mask'''] for example in examples]	)
    return {"pixel_values": pixel_values, "bool_masked_pos": mask}



def 		lowerCAmelCase__    (					)->  Optional[int]:
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    A__		 =					HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)	)
    if len(sys.argv	) == 2 and sys.argv[1].endswith('''.json'''	):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        A__ ,	A__ ,	A__		 =					parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]	)	)
    else:
        A__ ,	A__ ,	A__		 =					parser.parse_args_into_dataclasses()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry('''run_mim''' ,       UpperCamelCase_ ,       UpperCamelCase_	)

    # Setup logging
    logging.basicConfig(
        format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''' ,       datefmt='''%m/%d/%Y %H:%M:%S''' ,       handlers=[logging.StreamHandler(sys.stdout	)] ,       )

    if training_args.should_log:
        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
        transformers.utils.logging.set_verbosity_info()

    A__		 =					training_args.get_process_log_level()
    logger.setLevel(UpperCamelCase_	)
    transformers.utils.logging.set_verbosity(UpperCamelCase_	)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1	)}, 16-bits training: {training_args.fpaa}"	)
    logger.info(f"Training/evaluation parameters {training_args}"	)

    # Detecting last checkpoint.
    A__		 =					None
    if os.path.isdir(training_args.output_dir	) and training_args.do_train and not training_args.overwrite_output_dir:
        A__		 =					get_last_checkpoint(training_args.output_dir	)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir	)	) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                '''Use --overwrite_output_dir to overcome.'''	)
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                '''the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'''	)

    # Initialize our dataset.
    A__		 =					load_dataset(
        data_args.dataset_name ,       data_args.dataset_config_name ,       data_files=data_args.data_files ,       cache_dir=model_args.cache_dir ,       use_auth_token=True if model_args.use_auth_token else None ,       )

    # If we don't have a validation split, split off a percentage of train as validation.
    A__		 =					None if '''validation''' in ds.keys() else data_args.train_val_split
    if isinstance(data_args.train_val_split ,       UpperCamelCase_	) and data_args.train_val_split > 0.0:
        A__		 =					ds['''train'''].train_test_split(data_args.train_val_split	)
        A__		 =					split['''train''']
        A__		 =					split['''test''']

    # Create config
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    A__		 =					{
        '''cache_dir''': model_args.cache_dir,
        '''revision''': model_args.model_revision,
        '''use_auth_token''': True if model_args.use_auth_token else None,
    }
    if model_args.config_name_or_path:
        A__		 =					AutoConfig.from_pretrained(model_args.config_name_or_path ,       **UpperCamelCase_	)
    elif model_args.model_name_or_path:
        A__		 =					AutoConfig.from_pretrained(model_args.model_name_or_path ,       **UpperCamelCase_	)
    else:
        A__		 =					CONFIG_MAPPING[model_args.model_type]()
        logger.warning('''You are instantiating a new config instance from scratch.'''	)
        if model_args.config_overrides is not None:
            logger.info(f"Overriding config: {model_args.config_overrides}"	)
            config.update_from_string(model_args.config_overrides	)
            logger.info(f"New config: {config}"	)

    # make sure the decoder_type is "simmim" (only relevant for BEiT)
    if hasattr(UpperCamelCase_ ,       '''decoder_type'''	):
        A__		 =					'''simmim'''

    # adapt config
    A__		 =					model_args.image_size if model_args.image_size is not None else config.image_size
    A__		 =					model_args.patch_size if model_args.patch_size is not None else config.patch_size
    A__		 =					(
        model_args.encoder_stride if model_args.encoder_stride is not None else config.encoder_stride
    )

    config.update(
        {
            '''image_size''': model_args.image_size,
            '''patch_size''': model_args.patch_size,
            '''encoder_stride''': model_args.encoder_stride,
        }	)

    # create image processor
    if model_args.image_processor_name:
        A__		 =					AutoImageProcessor.from_pretrained(model_args.image_processor_name ,       **UpperCamelCase_	)
    elif model_args.model_name_or_path:
        A__		 =					AutoImageProcessor.from_pretrained(model_args.model_name_or_path ,       **UpperCamelCase_	)
    else:
        A__		 =					{
            conf.model_type: image_processor_class for conf, image_processor_class in IMAGE_PROCESSOR_MAPPING.items()
        }
        A__		 =					IMAGE_PROCESSOR_TYPES[model_args.model_type]()

    # create model
    if model_args.model_name_or_path:
        A__		 =					AutoModelForMaskedImageModeling.from_pretrained(
            model_args.model_name_or_path ,       from_tf=bool('''.ckpt''' in model_args.model_name_or_path	) ,       config=UpperCamelCase_ ,       cache_dir=model_args.cache_dir ,       revision=model_args.model_revision ,       use_auth_token=True if model_args.use_auth_token else None ,       )
    else:
        logger.info('''Training new model from scratch'''	)
        A__		 =					AutoModelForMaskedImageModeling.from_config(UpperCamelCase_	)

    if training_args.do_train:
        A__		 =					ds['''train'''].column_names
    else:
        A__		 =					ds['''validation'''].column_names

    if data_args.image_column_name is not None:
        A__		 =					data_args.image_column_name
    elif "image" in column_names:
        A__		 =					'''image'''
    elif "img" in column_names:
        A__		 =					'''img'''
    else:
        A__		 =					column_names[0]

    # transformations as done in original SimMIM paper
    # source: https://github.com/microsoft/SimMIM/blob/main/data/data_simmim.py
    A__		 =					Compose(
        [
            Lambda(lambda UpperCamelCase_	: img.convert('''RGB'''	) if img.mode != "RGB" else img	),
            RandomResizedCrop(model_args.image_size ,       scale=(0.67, 1.0) ,       ratio=(3.0 / 4.0, 4.0 / 3.0)	),
            RandomHorizontalFlip(),
            ToTensor(),
            Normalize(mean=image_processor.image_mean ,       std=image_processor.image_std	),
        ]	)

    # create mask generator
    A__		 =					MaskGenerator(
        input_size=model_args.image_size ,       mask_patch_size=data_args.mask_patch_size ,       model_patch_size=model_args.patch_size ,       mask_ratio=data_args.mask_ratio ,       )

    def preprocess_images(UpperCamelCase_				:       List[Any]	):
        A__		 =					[transforms(UpperCamelCase_	) for image in examples[image_column_name]]
        A__		 =					[mask_generator() for i in range(len(examples[image_column_name]	)	)]

        return examples

    if training_args.do_train:
        if "train" not in ds:
            raise ValueError('''--do_train requires a train dataset'''	)
        if data_args.max_train_samples is not None:
            A__		 =					ds['''train'''].shuffle(seed=training_args.seed	).select(range(data_args.max_train_samples	)	)
        # Set the training transforms
        ds["train"].set_transform(UpperCamelCase_	)

    if training_args.do_eval:
        if "validation" not in ds:
            raise ValueError('''--do_eval requires a validation dataset'''	)
        if data_args.max_eval_samples is not None:
            A__		 =					(
                ds['''validation'''].shuffle(seed=training_args.seed	).select(range(data_args.max_eval_samples	)	)
            )
        # Set the validation transforms
        ds["validation"].set_transform(UpperCamelCase_	)

    # Initialize our trainer
    A__		 =					Trainer(
        model=UpperCamelCase_ ,       args=UpperCamelCase_ ,       train_dataset=ds['''train'''] if training_args.do_train else None ,       eval_dataset=ds['''validation'''] if training_args.do_eval else None ,       tokenizer=UpperCamelCase_ ,       data_collator=UpperCamelCase_ ,       )

    # Training
    if training_args.do_train:
        A__		 =					None
        if training_args.resume_from_checkpoint is not None:
            A__		 =					training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            A__		 =					last_checkpoint
        A__		 =					trainer.train(resume_from_checkpoint=UpperCamelCase_	)
        trainer.save_model()
        trainer.log_metrics('''train''' ,       train_result.metrics	)
        trainer.save_metrics('''train''' ,       train_result.metrics	)
        trainer.save_state()

    # Evaluation
    if training_args.do_eval:
        A__		 =					trainer.evaluate()
        trainer.log_metrics('''eval''' ,       UpperCamelCase_	)
        trainer.save_metrics('''eval''' ,       UpperCamelCase_	)

    # Write model card and (optionally) push to hub
    A__		 =					{
        '''finetuned_from''': model_args.model_name_or_path,
        '''tasks''': '''masked-image-modeling''',
        '''dataset''': data_args.dataset_name,
        '''tags''': ['''masked-image-modeling'''],
    }
    if training_args.push_to_hub:
        trainer.push_to_hub(**UpperCamelCase_	)
    else:
        trainer.create_model_card(**UpperCamelCase_	)


if __name__ == "__main__":
   main()





from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import numpy as np
   import tensorflow as tf

   from transformers import (
       TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
       FlaubertConfig,
       TFFlaubertForMultipleChoice,
       TFFlaubertForQuestionAnsweringSimple,
       TFFlaubertForSequenceClassification,
       TFFlaubertForTokenClassification,
       TFFlaubertModel,
       TFFlaubertWithLMHeadModel,
   )



class     _UpperCAmelCase			:




     def __init__(				self ,						a__ ,						):
         A__		 =					parent
         A__		 =					1_3
         A__		 =					7
         A__		 =					True
         A__		 =					True
         A__		 =					True
         A__		 =					True
         A__		 =					True
         A__		 =					False
         A__		 =					False
         A__		 =					False
         A__		 =					2
         A__		 =					9_9
         A__		 =					0
         A__		 =					3_2
         A__		 =					2
         A__		 =					4
         A__		 =					0.1
         A__		 =					0.1
         A__		 =					5_1_2
         A__		 =					1_6
         A__		 =					2
         A__		 =					0.0_2
         A__		 =					3
         A__		 =					4
         A__		 =					'''last'''
         A__		 =					True
         A__		 =					None
         A__		 =					0




     def 		snake_case_		(				self):
         A__		 =					ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size)
         A__		 =					random_attention_mask([self.batch_size, self.seq_length] ,						dtype=tf.floataa)

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

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

         A__		 =					None
         A__		 =					None
         A__		 =					None
         if self.use_labels:
             A__		 =					ids_tensor([self.batch_size] ,						self.type_sequence_label_size)
             A__		 =					ids_tensor([self.batch_size, self.seq_length] ,						self.num_labels)
             A__		 =					ids_tensor([self.batch_size] ,						2 ,						dtype=tf.floataa)
             A__		 =					ids_tensor([self.batch_size] ,						self.num_choices)

         A__		 =					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 ,						bos_token_id=self.bos_token_id ,						)

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




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						):
         A__		 =					TFFlaubertModel(config=a__)
         A__		 =					{'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
         A__		 =					model(a__)

         A__		 =					[input_ids, input_mask]
         A__		 =					model(a__)
         self.parent.assertEqual(result.last_hidden_state.shape ,						(self.batch_size, self.seq_length, self.hidden_size))




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						):
         A__		 =					TFFlaubertWithLMHeadModel(a__)

         A__		 =					{'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
         A__		 =					model(a__)

         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size))




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						):
         A__		 =					TFFlaubertForQuestionAnsweringSimple(a__)

         A__		 =					{'''input_ids''': input_ids, '''lengths''': input_lengths}

         A__		 =					model(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 ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						):
         A__		 =					TFFlaubertForSequenceClassification(a__)

         A__		 =					{'''input_ids''': input_ids, '''lengths''': input_lengths}

         A__		 =					model(a__)

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




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						):
         A__		 =					self.num_labels
         A__		 =					TFFlaubertForTokenClassification(config=a__)
         A__		 =					{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         A__		 =					model(a__)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels))




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						):
         A__		 =					self.num_choices
         A__		 =					TFFlaubertForMultipleChoice(config=a__)
         A__		 =					tf.tile(tf.expand_dims(a__ ,						1) ,						(1, self.num_choices, 1))
         A__		 =					tf.tile(tf.expand_dims(a__ ,						1) ,						(1, self.num_choices, 1))
         A__		 =					tf.tile(tf.expand_dims(a__ ,						1) ,						(1, self.num_choices, 1))
         A__		 =					{
             '''input_ids''': multiple_choice_inputs_ids,
             '''attention_mask''': multiple_choice_input_mask,
             '''token_type_ids''': multiple_choice_token_type_ids,
         }
         A__		 =					model(a__)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.num_choices))




     def 		snake_case_		(				self):
         A__		 =					self.prepare_config_and_inputs()
         (
             (
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	
         )		 =					config_and_inputs
         A__		 =					{
             '''input_ids''': input_ids,
             '''token_type_ids''': token_type_ids,
             '''langs''': token_type_ids,
             '''lengths''': input_lengths,
         }
         return config, inputs_dict







@require_tf
class     _UpperCAmelCase			(       A__			,      A__			,      unittest.TestCase     ):
     UpperCamelCase__		    =				(
         (
             TFFlaubertModel,
             TFFlaubertWithLMHeadModel,
             TFFlaubertForSequenceClassification,
             TFFlaubertForQuestionAnsweringSimple,
             TFFlaubertForTokenClassification,
             TFFlaubertForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     UpperCamelCase__		    =				(
         (TFFlaubertWithLMHeadModel,) if is_tf_available() else ()
     )  # TODO (PVP): Check other models whether language generation is also applicable
     UpperCamelCase__		    =				(
         {
             '''feature-extraction''': TFFlaubertModel,
             '''fill-mask''': TFFlaubertWithLMHeadModel,
             '''question-answering''': TFFlaubertForQuestionAnsweringSimple,
             '''text-classification''': TFFlaubertForSequenceClassification,
             '''token-classification''': TFFlaubertForTokenClassification,
             '''zero-shot''': TFFlaubertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     UpperCamelCase__		    =				False
     UpperCamelCase__		    =				False




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__):
         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):
         A__		 =					TFFlaubertModelTester(self)
         A__		 =					ConfigTester(self ,						config_class=a__ ,						emb_dim=3_7)




     def 		snake_case_		(				self):
         self.config_tester.run_common_tests()




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_model(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_lm_head(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_qa(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_sequence_classif(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_for_token_classification(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_for_multiple_choice(*a__)




     @slow
     def 		snake_case_		(				self):
         for model_name in TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             A__		 =					TFFlaubertModel.from_pretrained(a__)
             self.assertIsNotNone(a__)







@require_tf
@require_sentencepiece
@require_tokenizers
class     _UpperCAmelCase			(       unittest.TestCase     ):




     @slow
     def 		snake_case_		(				self):
         A__		 =					TFFlaubertModel.from_pretrained('''jplu/tf-flaubert-small-cased''')

         A__		 =					tf.convert_to_tensor(
             [[0, 1_5_8, 7_3_5, 2_5_9_2, 1_4_2_4, 6_7_2_7, 8_2, 1]] ,						dtype=tf.intaa ,						)  # "J'aime flaubert !"

         A__		 =					model(a__)[0]
         A__		 =					tf.TensorShape((1, 8, 5_1_2))
         self.assertEqual(output.shape ,						a__)
         # compare the actual values for a slice.
         A__		 =					tf.convert_to_tensor(
             [
                 [
                     [-1.8_7_6_8_7_7_3, -1.5_6_6_5_5_5, 0.2_7_0_7_2_4_1_8],
                     [-1.6_9_2_0_0_3_8, -0.5_8_7_3_5_0_5, 1.9_3_2_9_5_9_9],
                     [-2.9_5_6_3_9_8_5, -1.6_9_9_3_8_3_5, 1.7_9_7_2_0_5_2],
                 ]
             ] ,						dtype=tf.floataa ,						)

         self.assertTrue(np.allclose(output[:, :3, :3].numpy() ,						expected_slice.numpy() ,						atol=1e-4))
