

import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


SCREAMING_SNAKE_CASE_        =      logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

SCREAMING_SNAKE_CASE_        =      {
    'base': AutoModel,
    'sequence-classification': AutoModelForSequenceClassification,
    'question-answering': AutoModelForQuestionAnswering,
    'pretraining': AutoModelForPreTraining,
    'token-classification': AutoModelForTokenClassification,
    'language-modeling': AutoModelWithLMHead,
    'summarization': AutoModelForSeqaSeqLM,
    'translation': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
SCREAMING_SNAKE_CASE_        =      {
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
SCREAMING_SNAKE_CASE_        =      sorted(arg_to_scheduler.keys())
SCREAMING_SNAKE_CASE_        =      '{' + ', '.join(arg_to_scheduler_choices) + '}'

class 				a  (	pl.LightningModule				):


						def __init__(       self						,					A_						,					A_=None						,					A_="base"						,					A_=None						,					A_=None						,					A_=None						,					**A_						,					):







										'''simple docstring'''


										super().__init__()
										# TODO: move to self.save_hyperparameters()
										# self.save_hyperparameters()
										# can also expand arguments into trainer signature for easier reading

										self.save_hyperparameters(A_	)
										_UpperCAmelCase       :		Any 		=    0
										_UpperCAmelCase       :		str 		=    Path(self.hparams.output_dir	)
										_UpperCAmelCase       :		List[str] 		=    self.hparams.cache_dir if self.hparams.cache_dir else None
										if config is None:
														_UpperCAmelCase       :		Union[str, Any] 		=    AutoConfig.from_pretrained(
														    self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path						,					**({"num_labels": num_labels} if num_labels is not None else {})						,					cache_dir=A_						,					**A_						,					)
										else:
														_UpperCAmelCase       :		PretrainedConfig 		=    config

										_UpperCAmelCase       :		List[Any] 		=    ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
										for p in extra_model_params:
														if getattr(self.hparams						,					A_						,					A_	):
																		assert hasattr(self.config						,					A_	), f'model config doesn\'t have a `{p}` attribute'
																		setattr(self.config						,					A_						,					getattr(self.hparams						,					A_	)	)

										if tokenizer is None:
														_UpperCAmelCase       :		int 		=    AutoTokenizer.from_pretrained(
														    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path						,					cache_dir=A_						,					)
										else:
														_UpperCAmelCase       :		PreTrainedTokenizer 		=    tokenizer
										_UpperCAmelCase       :		List[str] 		=    MODEL_MODES[mode]
										if model is None:
														_UpperCAmelCase       :		List[str] 		=    self.model_type.from_pretrained(
														    self.hparams.model_name_or_path						,					from_tf=bool(".ckpt" in self.hparams.model_name_or_path	)						,					config=self.config						,					cache_dir=A_						,					)
										else:
														_UpperCAmelCase       :		Tuple 		=    model


						def        _UpperCAmelCase				(       self						,					*A_						,					**A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    self.model_type.from_pretrained(*A_						,					**A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    arg_to_scheduler[self.hparams.lr_scheduler]
										_UpperCAmelCase       :		List[str] 		=    get_schedule_func(
										    self.opt						,					num_warmup_steps=self.hparams.warmup_steps						,					num_training_steps=self.total_steps()	)
										_UpperCAmelCase       :		Tuple 		=    {"scheduler": scheduler, "interval": "step", "frequency": 1}
										return scheduler


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    self.model
										_UpperCAmelCase       :		Dict 		=    ["bias", "LayerNorm.weight"]
										_UpperCAmelCase       :		Optional[int] 		=    [
										    {
										        "params": [
										            p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay	)
										        ],  # check this named paramters
										        "weight_decay": self.hparams.weight_decay,
										    },
										    {
										        "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay	)],
										        "weight_decay": 0.0,
										    },
										]
										if self.hparams.adafactor:
														_UpperCAmelCase       :		str 		=    Adafactor(
														    A_						,					lr=self.hparams.learning_rate						,					scale_parameter=A_						,					relative_step=A_	)

										else:
														_UpperCAmelCase       :		int 		=    AdamW(
														    A_						,					lr=self.hparams.learning_rate						,					eps=self.hparams.adam_epsilon	)
										_UpperCAmelCase       :		List[str] 		=    optimizer

										_UpperCAmelCase       :		Optional[int] 		=    self.get_lr_scheduler()

										return [optimizer], [scheduler]


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										return self.validation_step(A_						,					A_	)


						def        _UpperCAmelCase				(       self						,					A_	):







										'''simple docstring'''


										return self.validation_end(A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		str 		=    max(1						,					self.hparams.gpus	)  # TODO: consider num_tpu_cores
										_UpperCAmelCase       :		Tuple 		=    self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
										return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs


						def        _UpperCAmelCase				(       self						,					A_	):







										'''simple docstring'''


										if stage == "test":
														_UpperCAmelCase       :		Optional[Any] 		=    len(self.test_dataloader().dataset	)
										else:
														_UpperCAmelCase       :		Dict 		=    self.get_dataloader("train"						,					self.hparams.train_batch_size						,					shuffle=A_	)
														_UpperCAmelCase       :		Union[str, Any] 		=    len(self.train_dataloader().dataset	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_ = False	):







										'''simple docstring'''


										raise NotImplementedError("You must implement this for your task"	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return self.train_loader


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return self.get_dataloader("dev"						,					self.hparams.eval_batch_size						,					shuffle=A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return self.get_dataloader("test"						,					self.hparams.eval_batch_size						,					shuffle=A_	)


						def        _UpperCAmelCase				(       self						,					A_	):







										'''simple docstring'''


										return os.path.join(
										    self.hparams.data_dir						,					"cached_{}_{}_{}".format(
										        A_						,					list(filter(A_						,					self.hparams.model_name_or_path.split("/"	)	)	).pop()						,					str(self.hparams.max_seq_length	)						,					)						,					)


						@pl.utilities.rank_zero_only
						def        _UpperCAmelCase				(       self						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    self.output_dir.joinpath("best_tfmr"	)
										_UpperCAmelCase       :		int 		=    self.step_count
										self.model.save_pretrained(A_	)
										self.tokenizer.save_pretrained(A_	)





						@staticmethod
						def        _UpperCAmelCase				(       A_						,					A_	):







										'''simple docstring'''


										parser.add_argument(
										    "--model_name_or_path"						,					default=A_						,					type=A_						,					required=A_						,					help="Path to pretrained model or model identifier from huggingface.co/models"						,					)
										parser.add_argument(
										    "--config_name"						,					default=""						,					type=A_						,					help="Pretrained config name or path if not the same as model_name"	)
										parser.add_argument(
										    "--tokenizer_name"						,					default=A_						,					type=A_						,					help="Pretrained tokenizer name or path if not the same as model_name"						,					)
										parser.add_argument(
										    "--cache_dir"						,					default=str(Path(A_	).parent / "test_run" / "cache"	)						,					type=A_						,					help="Where do you want to store the pre-trained models downloaded from huggingface.co"						,					)
										parser.add_argument(
										    "--encoder_layerdrop"						,					type=A_						,					help="Encoder layer dropout probability (Optional). Goes into model.config"						,					)
										parser.add_argument(
										    "--decoder_layerdrop"						,					type=A_						,					help="Decoder layer dropout probability (Optional). Goes into model.config"						,					)
										parser.add_argument(
										    "--dropout"						,					type=A_						,					help="Dropout probability (Optional). Goes into model.config"						,					)
										parser.add_argument(
										    "--attention_dropout"						,					type=A_						,					help="Attention dropout probability (Optional). Goes into model.config"						,					)
										parser.add_argument("--learning_rate"						,					default=5e-5						,					type=A_						,					help="The initial learning rate for Adam."	)
										parser.add_argument(
										    "--lr_scheduler"						,					default="linear"						,					choices=A_						,					metavar=A_						,					type=A_						,					help="Learning rate scheduler"						,					)
										parser.add_argument("--weight_decay"						,					default=0.0						,					type=A_						,					help="Weight decay if we apply some."	)
										parser.add_argument("--adam_epsilon"						,					default=1e-8						,					type=A_						,					help="Epsilon for Adam optimizer."	)
										parser.add_argument("--warmup_steps"						,					default=0						,					type=A_						,					help="Linear warmup over warmup_steps."	)
										parser.add_argument("--num_workers"						,					default=4						,					type=A_						,					help="kwarg passed to DataLoader"	)
										parser.add_argument("--num_train_epochs"						,					dest="max_epochs"						,					default=3						,					type=A_	)
										parser.add_argument("--train_batch_size"						,					default=32						,					type=A_	)
										parser.add_argument("--eval_batch_size"						,					default=32						,					type=A_	)
										parser.add_argument("--adafactor"						,					action="store_true"	)







class 				a  (	pl.Callback				):


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										if (
										    trainer.is_global_zero and trainer.global_rank == 0
										):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
														pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.







class 				a  (	pl.Callback				):


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										for name, param in pl_module.model.rag.named_parameters():
														if param.grad is None:
																		print(A_	)







class 				a  (	pl.Callback				):


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    trainer.lr_schedulers[0]["scheduler"]
										_UpperCAmelCase       :		Optional[Any] 		=    {f'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr()	)}
										pl_module.logger.log_metrics(A_	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										rank_zero_info("***** Validation results *****"	)
										_UpperCAmelCase       :		Optional[int] 		=    trainer.callback_metrics
										# Log results
										for key in sorted(A_	):
														if key not in ["log", "progress_bar"]:
																		rank_zero_info("{} = {}\n".format(A_						,					str(metrics[key]	)	)	)





						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										rank_zero_info("***** Test results *****"	)
										_UpperCAmelCase       :		List[str] 		=    trainer.callback_metrics
										# Log and save results to file
										_UpperCAmelCase       :		Dict 		=    os.path.join(pl_module.hparams.output_dir						,					"test_results.txt"	)
										with open(A_						,					"w"	) as writer:
														for key in sorted(A_	):
																		if key not in ["log", "progress_bar"]:
																						rank_zero_info("{} = {}\n".format(A_						,					str(metrics[key]	)	)	)
																						writer.write("{} = {}\n".format(A_						,					str(metrics[key]	)	)	)







def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  str )   ->       None:
				#  To allow all pl args uncomment the following line
				#  parser = pl.Trainer.add_argparse_args(parser)
				parser.add_argument(
				    "--output_dir"  ,      default=str(Path(lowerCAmelCase ).parent / "test_run" / "model_checkpoints" )  ,      type=lowerCAmelCase  ,      help="The output directory where the model predictions and checkpoints will be written."  ,      )
				parser.add_argument(
				    "--fp16"  ,      action="store_true"  ,      help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"  ,      )

				parser.add_argument(
				    "--fp16_opt_level"  ,      type=lowerCAmelCase  ,      default="O2"  ,      help=(
				        "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
				        "See details at https://nvidia.github.io/apex/amp.html"
				    )  ,      )
				parser.add_argument("--n_tpu_cores"  ,      dest="tpu_cores"  ,      type=lowerCAmelCase )
				parser.add_argument("--max_grad_norm"  ,      dest="gradient_clip_val"  ,      default=1.0  ,      type=lowerCAmelCase  ,      help="Max gradient norm" )
				parser.add_argument("--do_train"  ,      action="store_true"  ,      help="Whether to run training." )
				parser.add_argument("--do_predict"  ,      action="store_true"  ,      help="Whether to run predictions on the test set." )
				parser.add_argument(
				    "--gradient_accumulation_steps"  ,      dest="accumulate_grad_batches"  ,      type=lowerCAmelCase  ,      default=1  ,      help="Number of updates steps to accumulate before performing a backward/update pass."  ,      )
				parser.add_argument("--seed"  ,      type=lowerCAmelCase  ,      default=42  ,      help="random seed for initialization" )
				parser.add_argument(
				    "--data_dir"  ,      default=str(Path(lowerCAmelCase ).parent / "test_run" / "dummy-train-data" )  ,      type=lowerCAmelCase  ,      help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."  ,      )





def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  BaseTransformer  ,      lowerCAmelCase:  argparse.Namespace  ,      lowerCAmelCase:  Any=None  ,      lowerCAmelCase:  Optional[Any]=True  ,      lowerCAmelCase:  Union[str, Any]=[]  ,      lowerCAmelCase:  Optional[Any]=None  ,      lowerCAmelCase:  Dict=None  ,      **lowerCAmelCase:  Any  ,      )   ->       Tuple:
				pl.seed_everything(args.seed )

				# init model
				_UpperCAmelCase       :		Optional[Any] 		=    Path(model.hparams.output_dir )
				odir.mkdir(exist_ok=lowerCAmelCase )

				# add custom checkpoints
				if checkpoint_callback is None:
								_UpperCAmelCase       :		Any 		=    pl.callbacks.ModelCheckpoint(
								    filepath=args.output_dir  ,      prefix="checkpoint"  ,      monitor="val_loss"  ,      mode="min"  ,      save_top_k=1 )
				if early_stopping_callback:
								extra_callbacks.append(lowerCAmelCase )
				if logging_callback is None:
								_UpperCAmelCase       :		Union[str, Any] 		=    LoggingCallback()

				_UpperCAmelCase       :		Any 		=    {}

				if args.fpaa:
								_UpperCAmelCase       :		Dict 		=    16

				if args.gpus > 1:
								_UpperCAmelCase       :		Tuple 		=    "auto"
								_UpperCAmelCase       :		List[str] 		=    "ddp"

				_UpperCAmelCase       :		Optional[Any] 		=    args.accumulate_grad_batches
				_UpperCAmelCase       :		int 		=    None
				_UpperCAmelCase       :		List[str] 		=    "auto"

				_UpperCAmelCase       :		Optional[int] 		=    pl.Trainer.from_argparse_args(
				    lowerCAmelCase  ,      weights_summary=lowerCAmelCase  ,      callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]  ,      logger=lowerCAmelCase  ,      val_check_interval=1  ,      num_sanity_val_steps=2  ,      **lowerCAmelCase  ,      )

				if args.do_train:
								trainer.fit(lowerCAmelCase )

				else:
								print("RAG modeling tests with new set functions successfuly executed!" )
				return trainer





import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch
					from torch import nn

					from transformers import (
					    MODEL_MAPPING,
					    BeitForImageClassification,
					    BeitForMaskedImageModeling,
					    BeitForSemanticSegmentation,
					    BeitModel,
					)
					from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					import PIL
					from PIL import Image

					from transformers import BeitImageProcessor

class 				a  :


						def __init__(       self						,					A_						,					A_=100						,					A_=13						,					A_=30						,					A_=2						,					A_=3						,					A_=True						,					A_=True						,					A_=32						,					A_=4						,					A_=4						,					A_=37						,					A_="gelu"						,					A_=0.1						,					A_=0.1						,					A_=10						,					A_=0.02						,					A_=3						,					A_=None						,					A_=[0, 1, 2, 3]						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    parent
										_UpperCAmelCase       :		Any 		=    100
										_UpperCAmelCase       :		Optional[Any] 		=    batch_size
										_UpperCAmelCase       :		Any 		=    image_size
										_UpperCAmelCase       :		int 		=    patch_size
										_UpperCAmelCase       :		Dict 		=    num_channels
										_UpperCAmelCase       :		int 		=    is_training
										_UpperCAmelCase       :		Union[str, Any] 		=    use_labels
										_UpperCAmelCase       :		Optional[Any] 		=    hidden_size
										_UpperCAmelCase       :		str 		=    num_hidden_layers
										_UpperCAmelCase       :		Tuple 		=    num_attention_heads
										_UpperCAmelCase       :		List[str] 		=    intermediate_size
										_UpperCAmelCase       :		List[str] 		=    hidden_act
										_UpperCAmelCase       :		List[Any] 		=    hidden_dropout_prob
										_UpperCAmelCase       :		Any 		=    attention_probs_dropout_prob
										_UpperCAmelCase       :		Dict 		=    type_sequence_label_size
										_UpperCAmelCase       :		Optional[Any] 		=    initializer_range
										_UpperCAmelCase       :		Optional[int] 		=    scope
										_UpperCAmelCase       :		Any 		=    out_indices
										_UpperCAmelCase       :		List[str] 		=    num_labels

										# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
										_UpperCAmelCase       :		Any 		=    (image_size // patch_size) ** 2
										_UpperCAmelCase       :		str 		=    num_patches + 1


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

										_UpperCAmelCase       :		Optional[Any] 		=    None
										_UpperCAmelCase       :		str 		=    None
										if self.use_labels:
														_UpperCAmelCase       :		Union[str, Any] 		=    ids_tensor([self.batch_size]						,					self.type_sequence_label_size	)
														_UpperCAmelCase       :		Optional[int] 		=    ids_tensor([self.batch_size, self.image_size, self.image_size]						,					self.num_labels	)

										_UpperCAmelCase       :		Tuple 		=    self.get_config()

										return config, pixel_values, labels, pixel_labels


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return BeitConfig(
										    vocab_size=self.vocab_size						,					image_size=self.image_size						,					patch_size=self.patch_size						,					num_channels=self.num_channels						,					hidden_size=self.hidden_size						,					num_hidden_layers=self.num_hidden_layers						,					num_attention_heads=self.num_attention_heads						,					intermediate_size=self.intermediate_size						,					hidden_act=self.hidden_act						,					hidden_dropout_prob=self.hidden_dropout_prob						,					attention_probs_dropout_prob=self.attention_probs_dropout_prob						,					is_decoder=A_						,					initializer_range=self.initializer_range						,					out_indices=self.out_indices						,					)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    BeitModel(config=A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		int 		=    model(A_	)
										self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, self.seq_length, self.hidden_size)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    BeitForMaskedImageModeling(config=A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		Any 		=    model(A_	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.seq_length - 1, self.vocab_size)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    self.type_sequence_label_size
										_UpperCAmelCase       :		Optional[int] 		=    BeitForImageClassification(A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		List[str] 		=    model(A_						,					labels=A_	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.type_sequence_label_size)	)

										# test greyscale images
										_UpperCAmelCase       :		Optional[Any] 		=    1
										_UpperCAmelCase       :		Tuple 		=    BeitForImageClassification(A_	)
										model.to(A_	)
										model.eval()

										_UpperCAmelCase       :		int 		=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
										_UpperCAmelCase       :		Dict 		=    model(A_						,					labels=A_	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.type_sequence_label_size)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    self.num_labels
										_UpperCAmelCase       :		Dict 		=    BeitForSemanticSegmentation(A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		str 		=    model(A_	)
										self.parent.assertEqual(
										    result.logits.shape						,					(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)	)
										_UpperCAmelCase       :		List[Any] 		=    model(A_						,					labels=A_	)
										self.parent.assertEqual(
										    result.logits.shape						,					(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)	)





						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    self.prepare_config_and_inputs()
										_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		Union[str, Any] 		=    config_and_inputs
										_UpperCAmelCase       :		int 		=    {"pixel_values": pixel_values}
										return config, inputs_dict







@require_torch
class 				a  (	UpperCAmelCase							, UpperCAmelCase							, unittest.TestCase				):
						_lowercase    						=			(
						    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
						    if is_torch_available()
						    else ()
						)
						_lowercase    						=			(
						    {
						        "feature-extraction": BeitModel,
						        "image-classification": BeitForImageClassification,
						        "image-segmentation": BeitForSemanticSegmentation,
						    }
						    if is_torch_available()
						    else {}
						)

						_lowercase    						=			False
						_lowercase    						=			False
						_lowercase    						=			False


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    BeitModelTester(self	)
										_UpperCAmelCase       :		List[Any] 		=    ConfigTester(self						,					config_class=A_						,					has_text_modality=A_						,					hidden_size=37	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										self.config_tester.run_common_tests()


						@unittest.skip(reason="BEiT does not use inputs_embeds"	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										pass


						@require_torch_multi_gpu
						@unittest.skip(reason="BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										pass


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		Dict 		=    self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														_UpperCAmelCase       :		Union[str, Any] 		=    model_class(A_	)
														self.assertIsInstance(model.get_input_embeddings()						,					(nn.Module)	)
														_UpperCAmelCase       :		Tuple 		=    model.get_output_embeddings()
														self.assertTrue(x is None or isinstance(A_						,					nn.Linear	)	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		int 		=    self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														_UpperCAmelCase       :		List[str] 		=    model_class(A_	)
														_UpperCAmelCase       :		Tuple 		=    inspect.signature(model.forward	)
														# signature.parameters is an OrderedDict => so arg_names order is deterministic
														_UpperCAmelCase       :		int 		=    [*signature.parameters.keys()]

														_UpperCAmelCase       :		int 		=    ["pixel_values"]
														self.assertListEqual(arg_names[:1]						,					A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_semantic_segmentation(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										if not self.model_tester.is_training:
														return

										_UpperCAmelCase		,							_UpperCAmelCase       :		Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs_for_common()
										_UpperCAmelCase       :		str 		=    True

										for model_class in self.all_model_classes:
														# we don't test BeitForMaskedImageModeling
														if model_class in [*get_values(A_	), BeitForMaskedImageModeling]:
																		continue

														_UpperCAmelCase       :		Union[str, Any] 		=    model_class(A_	)
														model.to(A_	)
														model.train()
														_UpperCAmelCase       :		str 		=    self._prepare_for_class(A_						,					A_						,					return_labels=A_	)
														_UpperCAmelCase       :		Optional[Any] 		=    model(**A_	).loss
														loss.backward()


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		Dict 		=    self.model_tester.prepare_config_and_inputs_for_common()
										if not self.model_tester.is_training:
														return

										_UpperCAmelCase       :		Tuple 		=    False
										_UpperCAmelCase       :		Any 		=    True

										for model_class in self.all_model_classes:
														# we don't test BeitForMaskedImageModeling
														if (
														    model_class in [*get_values(A_	), BeitForMaskedImageModeling]
														    or not model_class.supports_gradient_checkpointing
														):
																		continue

														_UpperCAmelCase       :		Optional[Any] 		=    model_class(A_	)
														model.gradient_checkpointing_enable()
														model.to(A_	)
														model.train()
														_UpperCAmelCase       :		Union[str, Any] 		=    self._prepare_for_class(A_						,					A_						,					return_labels=A_	)
														_UpperCAmelCase       :		Optional[Any] 		=    model(**A_	).loss
														loss.backward()


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		List[Any] 		=    self.model_tester.prepare_config_and_inputs_for_common()

										_UpperCAmelCase       :		Optional[Any] 		=    _config_zero_init(A_	)
										for model_class in self.all_model_classes:
														_UpperCAmelCase       :		Tuple 		=    model_class(config=A_	)
														for name, param in model.named_parameters():
																		# we skip lambda parameters as these require special initial values
																		# determined by config.layer_scale_init_value
																		if "lambda" in name:
																						continue
																		if param.requires_grad:
																						self.assertIn(
																						    ((param.data.mean() * 1e9).round() / 1e9).item()						,					[0.0, 1.0]						,					msg=f'Parameter {name} of model {model_class} seems not properly initialized'						,					)





						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_UpperCAmelCase       :		Dict 		=    BeitModel.from_pretrained(A_	)
														self.assertIsNotNone(A_	)







def 							__SCREAMING_SNAKE_CASE							(				)   ->       Any:
				_UpperCAmelCase       :		Optional[int] 		=    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png" )
				return image





@require_torch
@require_vision
class 				a  (	unittest.TestCase				):


						@cached_property
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"	) if is_vision_available() else None


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"	).to(A_	)

										_UpperCAmelCase       :		Dict 		=    self.default_image_processor
										_UpperCAmelCase       :		Optional[Any] 		=    prepare_img()
										_UpperCAmelCase       :		Any 		=    image_processor(images=A_						,					return_tensors="pt"	).pixel_values.to(A_	)

										# prepare bool_masked_pos
										_UpperCAmelCase       :		Dict 		=    torch.ones((1, 196)						,					dtype=torch.bool	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		List[str] 		=    model(pixel_values=A_						,					bool_masked_pos=A_	)
										_UpperCAmelCase       :		Optional[int] 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		Optional[int] 		=    torch.Size((1, 196, 8192)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		Tuple 		=    torch.tensor(
										    [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]	).to(A_	)

										self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]						,					A_						,					atol=1e-2	)	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"	).to(A_	)

										_UpperCAmelCase       :		str 		=    self.default_image_processor
										_UpperCAmelCase       :		Union[str, Any] 		=    prepare_img()
										_UpperCAmelCase       :		List[Any] 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		Tuple 		=    model(**A_	)
										_UpperCAmelCase       :		Dict 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		List[str] 		=    torch.Size((1, 1000)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		int 		=    torch.tensor([-1.23_85, -1.09_87, -1.01_08]	).to(A_	)

										self.assertTrue(torch.allclose(logits[0, :3]						,					A_						,					atol=1e-4	)	)

										_UpperCAmelCase       :		List[Any] 		=    281
										self.assertEqual(logits.argmax(-1	).item()						,					A_	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    BeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"	).to(
										    A_	)

										_UpperCAmelCase       :		str 		=    self.default_image_processor
										_UpperCAmelCase       :		List[str] 		=    prepare_img()
										_UpperCAmelCase       :		List[Any] 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		Any 		=    model(**A_	)
										_UpperCAmelCase       :		List[Any] 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		List[str] 		=    torch.Size((1, 21841)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		Union[str, Any] 		=    torch.tensor([1.68_81, -0.27_87, 0.59_01]	).to(A_	)

										self.assertTrue(torch.allclose(logits[0, :3]						,					A_						,					atol=1e-4	)	)

										_UpperCAmelCase       :		List[str] 		=    2396
										self.assertEqual(logits.argmax(-1	).item()						,					A_	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"	)
										_UpperCAmelCase       :		str 		=    model.to(A_	)

										_UpperCAmelCase       :		int 		=    BeitImageProcessor(do_resize=A_						,					size=640						,					do_center_crop=A_	)

										_UpperCAmelCase       :		Union[str, Any] 		=    load_dataset("hf-internal-testing/fixtures_ade20k"						,					split="test"	)
										_UpperCAmelCase       :		Any 		=    Image.open(ds[0]["file"]	)
										_UpperCAmelCase       :		int 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		Dict 		=    model(**A_	)
										_UpperCAmelCase       :		Optional[int] 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		str 		=    torch.Size((1, 150, 160, 160)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		str 		=    version.parse(PIL.__version__	) < version.parse("9.0.0"	)

										if is_pillow_less_than_a:
														_UpperCAmelCase       :		int 		=    torch.tensor(
														    [
														        [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
														        [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
														        [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
														    ]						,					device=A_						,					)
										else:
														_UpperCAmelCase       :		Any 		=    torch.tensor(
														    [
														        [[-4.89_60, -2.36_88, -3.03_55], [-2.84_78, -0.98_36, -1.74_18], [-2.94_49, -1.33_32, -2.14_56]],
														        [[-5.80_81, -3.41_24, -4.10_06], [-3.85_61, -2.20_81, -3.03_23], [-3.83_65, -2.46_01, -3.36_69]],
														        [[-0.03_09, 3.98_68, 4.05_40], [2.96_40, 4.68_77, 4.99_76], [3.20_81, 4.76_90, 4.99_42]],
														    ]						,					device=A_						,					)

										self.assertTrue(torch.allclose(logits[0, :3, :3, :3]						,					A_						,					atol=1e-4	)	)





						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"	)
										_UpperCAmelCase       :		Dict 		=    model.to(A_	)

										_UpperCAmelCase       :		Dict 		=    BeitImageProcessor(do_resize=A_						,					size=640						,					do_center_crop=A_	)

										_UpperCAmelCase       :		Optional[Any] 		=    load_dataset("hf-internal-testing/fixtures_ade20k"						,					split="test"	)
										_UpperCAmelCase       :		List[str] 		=    Image.open(ds[0]["file"]	)
										_UpperCAmelCase       :		Union[str, Any] 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		List[Any] 		=    model(**A_	)

										_UpperCAmelCase       :		Tuple 		=    outputs.logits.detach().cpu()

										_UpperCAmelCase       :		str 		=    image_processor.post_process_semantic_segmentation(outputs=A_						,					target_sizes=[(500, 300)]	)
										_UpperCAmelCase       :		Dict 		=    torch.Size((500, 300)	)
										self.assertEqual(segmentation[0].shape						,					A_	)

										_UpperCAmelCase       :		Union[str, Any] 		=    image_processor.post_process_semantic_segmentation(outputs=A_	)
										_UpperCAmelCase       :		Union[str, Any] 		=    torch.Size((160, 160)	)
										self.assertEqual(segmentation[0].shape						,					A_	)


