






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


_a   			=						logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

_a   			=						{
    '''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
_a   			=						{
    '''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
}
_a   			=						sorted(arg_to_scheduler.keys())
_a   			=						'''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''





class __lowerCamelCase				(   pl.LightningModule):





	"""simple docstring"""






	def __init__(							self ,					UpperCAmelCase ,					UpperCAmelCase=None ,					UpperCAmelCase="base" ,					UpperCAmelCase=None ,					UpperCAmelCase=None ,					UpperCAmelCase=None ,					**UpperCAmelCase ,					):
							"""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(UpperCAmelCase					)
							_UpperCAmelCase             =    0
							_UpperCAmelCase             =    Path(self.hparams.output_dir					)
							_UpperCAmelCase             =    self.hparams.cache_dir if self.hparams.cache_dir else None
							if config is None:
													_UpperCAmelCase             =    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=UpperCAmelCase ,					**UpperCAmelCase ,					)
							else:
													_UpperCAmelCase             =    config

							_UpperCAmelCase             =    ('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
							for p in extra_model_params:
													if getattr(self.hparams ,					UpperCAmelCase ,					UpperCAmelCase					):
																			assert hasattr(self.config ,					UpperCAmelCase					), F"""model config doesn't have a `{p}` attribute"""
																			setattr(self.config ,					UpperCAmelCase ,					getattr(self.hparams ,					UpperCAmelCase					)					)

							if tokenizer is None:
													_UpperCAmelCase             =    AutoTokenizer.from_pretrained(
													    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path ,					cache_dir=UpperCAmelCase ,					)
							else:
													_UpperCAmelCase             =    tokenizer
							_UpperCAmelCase             =    MODEL_MODES[mode]
							if model is None:
													_UpperCAmelCase             =    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=UpperCAmelCase ,					)
							else:
													_UpperCAmelCase             =    model

	def  UpperCamelCase					(							self ,					*UpperCAmelCase ,					**UpperCAmelCase					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.model_type.from_pretrained(*UpperCAmelCase ,					**UpperCAmelCase					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







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

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.model
							_UpperCAmelCase             =    ['bias', 'LayerNorm.weight']
							_UpperCAmelCase             =    [
							    {
							        'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay					)],
							        '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             =    Adafactor(
													    UpperCAmelCase ,					lr=self.hparams.learning_rate ,					scale_parameter=UpperCAmelCase ,					relative_step=UpperCAmelCase					)

							else:
													_UpperCAmelCase             =    AdamW(
													    UpperCAmelCase ,					lr=self.hparams.learning_rate ,					eps=self.hparams.adam_epsilon					)
							_UpperCAmelCase             =    optimizer

							_UpperCAmelCase             =    self.get_lr_scheduler()

							return [optimizer], [scheduler]

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







							return self.validation_step(UpperCAmelCase ,					UpperCAmelCase					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase					):
							"""simple docstring"""







							return self.validation_end(UpperCAmelCase					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    max(1 ,					self.hparams.gpus					)  # TODO: consider num_tpu_cores
							_UpperCAmelCase             =    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 ,					UpperCAmelCase					):
							"""simple docstring"""







							if stage == "test":
													_UpperCAmelCase             =    len(self.test_dataloader().dataset					)
							else:
													_UpperCAmelCase             =    self.get_dataloader('train' ,					self.hparams.train_batch_size ,					shuffle=UpperCAmelCase					)
													_UpperCAmelCase             =    len(self.train_dataloader().dataset					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase = 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=UpperCAmelCase					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return self.get_dataloader('test' ,					self.hparams.eval_batch_size ,					shuffle=UpperCAmelCase					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase					):
							"""simple docstring"""







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

	@pl.utilities.rank_zero_only
	def  UpperCamelCase					(							self ,					UpperCAmelCase					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.output_dir.joinpath('best_tfmr'					)
							_UpperCAmelCase             =    self.step_count
							self.model.save_pretrained(UpperCAmelCase					)
							self.tokenizer.save_pretrained(UpperCAmelCase					)

	@staticmethod
	def  UpperCamelCase					(							UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







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





class __lowerCamelCase				(   pl.Callback):





	"""simple docstring"""






	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""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 __lowerCamelCase				(   pl.Callback):





	"""simple docstring"""






	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







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

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







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

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







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



def __A					(   __lowerCAmelCase				,	__lowerCAmelCase			)->						None:





						"""simple docstring"""



						parser.add_argument(
						    '--output_dir'				,	default=__lowerCAmelCase				,	type=__lowerCAmelCase				,	required=__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=__lowerCAmelCase				,	type=__lowerCAmelCase				,	required=__lowerCAmelCase				,	help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'				,	)


def __A					(   __lowerCAmelCase				,	__lowerCAmelCase				,	__lowerCAmelCase=None				,	__lowerCAmelCase=True				,	__lowerCAmelCase=None				,	__lowerCAmelCase=[]				,	__lowerCAmelCase=None				,	__lowerCAmelCase=None				,	**__lowerCAmelCase				,	)->						Tuple:





						"""simple docstring"""



						pl.seed_everything(args.seed			)

						# init model
						_UpperCAmelCase             =    Path(model.hparams.output_dir			)
						odir.mkdir(exist_ok=__lowerCAmelCase			)

						# add custom checkpoints
						if checkpoint_callback is None:
												_UpperCAmelCase             =    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             =    LoggingCallback()

						_UpperCAmelCase             =    {}

						# TODO: remove with PyTorch 1.6 since pl uses native amp
						if args.fpaa:
												_UpperCAmelCase             =    16
												# train_params["amp_level"] = args.fp16_opt_level

						if args.gpus > 1:
												_UpperCAmelCase             =    'auto'  # "ddp"
												_UpperCAmelCase             =    'ddp'

						_UpperCAmelCase             =    args.accumulate_grad_batches
						_UpperCAmelCase             =    None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
						_UpperCAmelCase             =    'auto'

						_UpperCAmelCase             =    pl.Trainer.from_argparse_args(
						    __lowerCAmelCase				,	weights_summary=__lowerCAmelCase				,	callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]				,	logger=__lowerCAmelCase				,	**__lowerCAmelCase				,	)

						if args.do_train:
												trainer.fit(__lowerCAmelCase			)

						return trainer








from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import SwinConfig
from transformers.testing_utils import require_tf, require_vision, slow, to_atuple
from transformers.utils import cached_property, is_tf_available, is_vision_available

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


if is_tf_available():
		import tensorflow as tf

		from transformers.models.swin.modeling_tf_swin import (
		    TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFSwinForImageClassification,
		    TFSwinForMaskedImageModeling,
		    TFSwinModel,
		)


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor





class __lowerCamelCase				:





	"""simple docstring"""






	def __init__(							self ,					UpperCAmelCase ,					UpperCAmelCase=13 ,					UpperCAmelCase=32 ,					UpperCAmelCase=2 ,					UpperCAmelCase=3 ,					UpperCAmelCase=16 ,					UpperCAmelCase=[1, 2, 1] ,					UpperCAmelCase=[2, 2, 4] ,					UpperCAmelCase=2 ,					UpperCAmelCase=2.0 ,					UpperCAmelCase=True ,					UpperCAmelCase=0.0 ,					UpperCAmelCase=0.0 ,					UpperCAmelCase=0.1 ,					UpperCAmelCase="gelu" ,					UpperCAmelCase=False ,					UpperCAmelCase=True ,					UpperCAmelCase=0.02 ,					UpperCAmelCase=1e-5 ,					UpperCAmelCase=True ,					UpperCAmelCase=None ,					UpperCAmelCase=True ,					UpperCAmelCase=10 ,					UpperCAmelCase=8 ,					):
							"""simple docstring"""







							_UpperCAmelCase             =    parent
							_UpperCAmelCase             =    batch_size
							_UpperCAmelCase             =    image_size
							_UpperCAmelCase             =    patch_size
							_UpperCAmelCase             =    num_channels
							_UpperCAmelCase             =    embed_dim
							_UpperCAmelCase             =    depths
							_UpperCAmelCase             =    num_heads
							_UpperCAmelCase             =    window_size
							_UpperCAmelCase             =    mlp_ratio
							_UpperCAmelCase             =    qkv_bias
							_UpperCAmelCase             =    hidden_dropout_prob
							_UpperCAmelCase             =    attention_probs_dropout_prob
							_UpperCAmelCase             =    drop_path_rate
							_UpperCAmelCase             =    hidden_act
							_UpperCAmelCase             =    use_absolute_embeddings
							_UpperCAmelCase             =    patch_norm
							_UpperCAmelCase             =    layer_norm_eps
							_UpperCAmelCase             =    initializer_range
							_UpperCAmelCase             =    is_training
							_UpperCAmelCase             =    scope
							_UpperCAmelCase             =    use_labels
							_UpperCAmelCase             =    type_sequence_label_size
							_UpperCAmelCase             =    encoder_stride

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







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

							_UpperCAmelCase             =    None
							if self.use_labels:
													_UpperCAmelCase             =    ids_tensor([self.batch_size] ,					self.type_sequence_label_size					)

							_UpperCAmelCase             =    self.get_config()

							return config, pixel_values, labels

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return SwinConfig(
							    image_size=self.image_size ,					patch_size=self.patch_size ,					num_channels=self.num_channels ,					embed_dim=self.embed_dim ,					depths=self.depths ,					num_heads=self.num_heads ,					window_size=self.window_size ,					mlp_ratio=self.mlp_ratio ,					qkv_bias=self.qkv_bias ,					hidden_dropout_prob=self.hidden_dropout_prob ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob ,					drop_path_rate=self.drop_path_rate ,					hidden_act=self.hidden_act ,					use_absolute_embeddings=self.use_absolute_embeddings ,					path_norm=self.patch_norm ,					layer_norm_eps=self.layer_norm_eps ,					initializer_range=self.initializer_range ,					encoder_stride=self.encoder_stride ,					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







							_UpperCAmelCase             =    TFSwinModel(config=UpperCAmelCase					)
							_UpperCAmelCase             =    model(UpperCAmelCase					)

							_UpperCAmelCase             =    ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths					) - 1))
							_UpperCAmelCase             =    int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

							self.parent.assertEqual(result.last_hidden_state.shape ,					(self.batch_size, expected_seq_len, expected_dim)					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







							_UpperCAmelCase             =    TFSwinForMaskedImageModeling(config=UpperCAmelCase					)
							_UpperCAmelCase             =    model(UpperCAmelCase					)
							self.parent.assertEqual(
							    result.logits.shape ,					(self.batch_size, self.num_channels, self.image_size, self.image_size)					)

							# test greyscale images
							_UpperCAmelCase             =    1
							_UpperCAmelCase             =    TFSwinForMaskedImageModeling(UpperCAmelCase					)

							_UpperCAmelCase             =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
							_UpperCAmelCase             =    model(UpperCAmelCase					)
							self.parent.assertEqual(result.logits.shape ,					(self.batch_size, 1, self.image_size, self.image_size)					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.type_sequence_label_size
							_UpperCAmelCase             =    TFSwinForImageClassification(UpperCAmelCase					)
							_UpperCAmelCase             =    model(UpperCAmelCase ,					labels=UpperCAmelCase					)
							self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.type_sequence_label_size)					)

							# test greyscale images
							_UpperCAmelCase             =    1
							_UpperCAmelCase             =    TFSwinForImageClassification(UpperCAmelCase					)
							_UpperCAmelCase             =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
							_UpperCAmelCase             =    model(UpperCAmelCase					)
							self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.type_sequence_label_size)					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.prepare_config_and_inputs()
							_UpperCAmelCase      ,				_UpperCAmelCase      ,				_UpperCAmelCase             =    config_and_inputs
							_UpperCAmelCase             =    {'pixel_values': pixel_values}
							return config, inputs_dict





@require_tf
class __lowerCamelCase				(   snake_case__  , snake_case__  , unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__			=   (
	    (
	        TFSwinModel,
	        TFSwinForImageClassification,
	        TFSwinForMaskedImageModeling,
	    )
	    if is_tf_available()
	    else ()
	)
	UpperCamelCase__			=   (
	    {"feature-extraction": TFSwinModel, "image-classification": TFSwinForImageClassification}
	    if is_tf_available()
	    else {}
	)

	UpperCamelCase__			=   False
	UpperCamelCase__			=   False
	UpperCamelCase__			=   False
	UpperCamelCase__			=   False

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    TFSwinModelTester(self					)
							_UpperCAmelCase             =    ConfigTester(self ,					config_class=UpperCAmelCase ,					embed_dim=37					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							self.create_and_test_config_common_properties()
							self.config_tester.create_and_test_config_to_json_string()
							self.config_tester.create_and_test_config_to_json_file()
							self.config_tester.create_and_test_config_from_and_save_pretrained()
							self.config_tester.create_and_test_config_with_num_labels()
							self.config_tester.check_config_can_be_init_without_params()
							self.config_tester.check_config_arguments_init()

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCAmelCase					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase					)

	@unittest.skip(reason='Swin does not use inputs_embeds'					)
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							pass

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase      ,				_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													_UpperCAmelCase             =    model_class(UpperCAmelCase					)
													self.assertIsInstance(model.get_input_embeddings() ,					tf.keras.layers.Layer					)
													_UpperCAmelCase             =    model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(UpperCAmelCase ,					tf.keras.layers.Dense					)					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase      ,				_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													_UpperCAmelCase             =    model_class(UpperCAmelCase					)
													_UpperCAmelCase             =    inspect.signature(model.call					)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													_UpperCAmelCase             =    [*signature.parameters.keys()]

													_UpperCAmelCase             =    ['pixel_values']
													self.assertListEqual(arg_names[:1] ,					UpperCAmelCase					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase      ,				_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs_for_common()
							_UpperCAmelCase             =    True

							for model_class in self.all_model_classes:
													_UpperCAmelCase             =    True
													_UpperCAmelCase             =    False
													_UpperCAmelCase             =    True
													_UpperCAmelCase             =    model_class(UpperCAmelCase					)
													_UpperCAmelCase             =    model(**self._prepare_for_class(UpperCAmelCase ,					UpperCAmelCase					)					)
													_UpperCAmelCase             =    outputs.attentions
													_UpperCAmelCase             =    len(self.model_tester.depths					)
													self.assertEqual(len(UpperCAmelCase					) ,					UpperCAmelCase					)

													# check that output_attentions also work using config
													del inputs_dict["output_attentions"]
													_UpperCAmelCase             =    True
													_UpperCAmelCase             =    config.window_size**2
													_UpperCAmelCase             =    model_class(UpperCAmelCase					)
													_UpperCAmelCase             =    model(**self._prepare_for_class(UpperCAmelCase ,					UpperCAmelCase					)					)
													_UpperCAmelCase             =    outputs.attentions
													self.assertEqual(len(UpperCAmelCase					) ,					UpperCAmelCase					)

													self.assertListEqual(
													    list(attentions[0].shape[-3:]					) ,					[self.model_tester.num_heads[0], window_size_squared, window_size_squared] ,					)
													_UpperCAmelCase             =    len(UpperCAmelCase					)

													# Check attention is always last and order is fine
													_UpperCAmelCase             =    True
													_UpperCAmelCase             =    True
													_UpperCAmelCase             =    model_class(UpperCAmelCase					)
													_UpperCAmelCase             =    model(**self._prepare_for_class(UpperCAmelCase ,					UpperCAmelCase					)					)

													if hasattr(self.model_tester ,					'num_hidden_states_types'					):
																			_UpperCAmelCase             =    self.model_tester.num_hidden_states_types
													else:
																			# also another +1 for reshaped_hidden_states
																			_UpperCAmelCase             =    2
													self.assertEqual(out_len + added_hidden_states ,					len(UpperCAmelCase					)					)

													_UpperCAmelCase             =    outputs.attentions

													self.assertEqual(len(UpperCAmelCase					) ,					UpperCAmelCase					)

													self.assertListEqual(
													    list(self_attentions[0].shape[-3:]					) ,					[self.model_tester.num_heads[0], window_size_squared, window_size_squared] ,					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase					):
							"""simple docstring"""







							_UpperCAmelCase             =    model_class(UpperCAmelCase					)
							_UpperCAmelCase             =    model(**self._prepare_for_class(UpperCAmelCase ,					UpperCAmelCase					)					)
							_UpperCAmelCase             =    outputs.hidden_states

							_UpperCAmelCase             =    getattr(
							    self.model_tester ,					'expected_num_hidden_layers' ,					len(self.model_tester.depths					) + 1					)
							self.assertEqual(len(UpperCAmelCase					) ,					UpperCAmelCase					)

							# Swin has a different seq_length
							_UpperCAmelCase             =    to_atuple(config.patch_size					)

							_UpperCAmelCase             =    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

							self.assertListEqual(
							    list(hidden_states[0].shape[-2:]					) ,					[num_patches, self.model_tester.embed_dim] ,					)

							_UpperCAmelCase             =    outputs.reshaped_hidden_states
							self.assertEqual(len(UpperCAmelCase					) ,					UpperCAmelCase					)

							_UpperCAmelCase      ,				_UpperCAmelCase      ,				_UpperCAmelCase      ,				_UpperCAmelCase             =    reshaped_hidden_states[0].shape

							_UpperCAmelCase             =    tf.reshape(reshaped_hidden_states[0] ,					(batch_size, num_channels, height * width)					)
							_UpperCAmelCase             =    tf.transpose(UpperCAmelCase ,					(0, 2, 1)					)

							self.assertListEqual(
							    list(reshaped_hidden_states.shape[-2:]					) ,					[num_patches, self.model_tester.embed_dim] ,					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase      ,				_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs_for_common()

							_UpperCAmelCase             =    to_atuple(self.model_tester.image_size					)

							for model_class in self.all_model_classes:
													_UpperCAmelCase             =    True
													self.check_hidden_states_output(UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase					)

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													_UpperCAmelCase             =    True

													self.check_hidden_states_output(UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase					)

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase      ,				_UpperCAmelCase             =    self.model_tester.prepare_config_and_inputs_for_common()
							_UpperCAmelCase             =    3

							_UpperCAmelCase             =    to_atuple(self.model_tester.image_size					)
							_UpperCAmelCase             =    to_atuple(config.patch_size					)

							_UpperCAmelCase             =    image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
							_UpperCAmelCase             =    image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

							for model_class in self.all_model_classes:
													_UpperCAmelCase             =    True
													self.check_hidden_states_output(UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase ,					(padded_height, padded_width)					)

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													_UpperCAmelCase             =    True
													self.check_hidden_states_output(UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase ,					(padded_height, padded_width)					)

	@slow
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							for model_name in TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_UpperCAmelCase             =    TFSwinModel.from_pretrained(UpperCAmelCase					)
													self.assertIsNotNone(UpperCAmelCase					)





@require_vision
@require_tf
class __lowerCamelCase				(   unittest.TestCase):





	"""simple docstring"""






	@cached_property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return (
							    AutoImageProcessor.from_pretrained('microsoft/swin-tiny-patch4-window7-224'					)
							    if is_vision_available()
							    else None
							)

	@slow
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    TFSwinForImageClassification.from_pretrained('microsoft/swin-tiny-patch4-window7-224'					)
							_UpperCAmelCase             =    self.default_image_processor

							_UpperCAmelCase             =    Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'					)
							_UpperCAmelCase             =    image_processor(images=UpperCAmelCase ,					return_tensors='tf'					)

							# forward pass
							_UpperCAmelCase             =    model(UpperCAmelCase					)

							# verify the logits
							_UpperCAmelCase             =    tf.TensorShape((1, 1000)					)
							self.assertEqual(outputs.logits.shape ,					UpperCAmelCase					)
							_UpperCAmelCase             =    tf.constant([-0.09_48, -0.64_54, -0.09_21]					)
							self.assertTrue(np.allclose(outputs.logits[0, :3] ,					UpperCAmelCase ,					atol=1e-4					)					)
