






import copy
import inspect
import unittest

from transformers import PretrainedConfig, SwiftFormerConfig
from transformers.testing_utils import (
    require_torch,
    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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch
				from torch import nn

				from transformers import SwiftFormerForImageClassification, SwiftFormerModel
				from transformers.models.swiftformer.modeling_swiftformer import SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from PIL import Image

				from transformers import ViTImageProcessor

class 		_snake_case  :






			def __init__(       self		,					a		,					a=13		,					a=3		,					a=True		,					a=True		,					a=0.1		,					a=0.1		,					a=224		,					a=1000		,					a=[3, 3, 6, 4]		,					a=[48, 56, 112, 220]		,					)      ->       Tuple:
						SCREAMING_SNAKE_CASE						= parent
						SCREAMING_SNAKE_CASE						= batch_size
						SCREAMING_SNAKE_CASE						= num_channels
						SCREAMING_SNAKE_CASE						= is_training
						SCREAMING_SNAKE_CASE						= use_labels
						SCREAMING_SNAKE_CASE						= hidden_dropout_prob
						SCREAMING_SNAKE_CASE						= attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE						= num_labels
						SCREAMING_SNAKE_CASE						= image_size
						SCREAMING_SNAKE_CASE						= layer_depths
						SCREAMING_SNAKE_CASE						= embed_dims






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Any:
						SCREAMING_SNAKE_CASE						= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

						SCREAMING_SNAKE_CASE						= None
						if self.use_labels:
									SCREAMING_SNAKE_CASE						= ids_tensor([self.batch_size]		,					self.num_labels)

						SCREAMING_SNAKE_CASE						= self.get_config()

						return config, pixel_values, labels






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						return SwiftFormerConfig(
						    depths=self.layer_depths		,					embed_dims=self.embed_dims		,					mlp_ratio=4		,					downsamples=[True, True, True, True]		,					hidden_act='gelu'		,					num_labels=self.num_labels		,					down_patch_size=3		,					down_stride=2		,					down_pad=1		,					drop_rate=0.0		,					drop_path_rate=0.0		,					use_layer_scale=a		,					layer_scale_init_value=1E-5		,					)






			def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a)      ->       Any:
						SCREAMING_SNAKE_CASE						= SwiftFormerModel(config=a)
						model.to(a)
						model.eval()
						SCREAMING_SNAKE_CASE						= model(a)
						self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.embed_dims[-1], 7, 7))






			def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a)      ->       int:
						SCREAMING_SNAKE_CASE						= self.num_labels
						SCREAMING_SNAKE_CASE						= SwiftFormerForImageClassification(a)
						model.to(a)
						model.eval()
						SCREAMING_SNAKE_CASE						= model(a		,					labels=a)
						self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_labels))

						SCREAMING_SNAKE_CASE						= SwiftFormerForImageClassification(a)
						model.to(a)
						model.eval()

						SCREAMING_SNAKE_CASE						= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])
						SCREAMING_SNAKE_CASE						= model(a)
						self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_labels))







			def    SCREAMING_SNAKE_CASE__    (       self)      ->       int:
						((SCREAMING_SNAKE_CASE)				,						(SCREAMING_SNAKE_CASE)				,						(SCREAMING_SNAKE_CASE))						= self.prepare_config_and_inputs()
						SCREAMING_SNAKE_CASE						= {'pixel_values': pixel_values}
						return config, inputs_dict



@require_torch
class 		_snake_case  (  A__   ,    A__   ,    unittest.TestCase					):
			_lowercase		:    Dict			     =							(SwiftFormerModel, SwiftFormerForImageClassification) if is_torch_available() else ()
			_lowercase		:    Optional[int]			     =							(
			    {'''feature-extraction''': SwiftFormerModel, '''image-classification''': SwiftFormerForImageClassification}
			    if is_torch_available()
			    else {}
			)

			_lowercase		:    Dict			     =							False
			_lowercase		:    List[Any]			     =							False
			_lowercase		:    Optional[int]			     =							False
			_lowercase		:    Dict			     =							False
			_lowercase		:    Union[str, Any]			     =							False






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						SCREAMING_SNAKE_CASE						= SwiftFormerModelTester(self)
						SCREAMING_SNAKE_CASE						= ConfigTester(
						    self		,					config_class=a		,					has_text_modality=a		,					hidden_size=37		,					num_attention_heads=12		,					num_hidden_layers=12		,					)






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						self.config_tester.run_common_tests()






			@unittest.skip(reason='SwiftFormer does not use inputs_embeds')
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						pass






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Dict:
						SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE						= model_class(a)
									SCREAMING_SNAKE_CASE						= model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(a		,					nn.Linear))






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs_for_common()

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

									SCREAMING_SNAKE_CASE						= ['pixel_values']
									self.assertListEqual(arg_names[:1]		,					a)






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*a)






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*a)






			@slow
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Dict:
						for model_name in SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									SCREAMING_SNAKE_CASE						= SwiftFormerModel.from_pretrained(a)
									self.assertIsNotNone(a)






			@unittest.skip(reason='SwiftFormer does not output attentions')
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Any:
						pass






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						def check_hidden_states_output(a		,					a		,					a):
									SCREAMING_SNAKE_CASE						= model_class(a)
									model.to(a)
									model.eval()

									with torch.no_grad():
												SCREAMING_SNAKE_CASE						= model(**self._prepare_for_class(a		,					a))

									SCREAMING_SNAKE_CASE						= outputs.hidden_states

									SCREAMING_SNAKE_CASE						= 8
									self.assertEqual(len(a)		,					a)  # TODO

									# SwiftFormer's feature maps are of shape (batch_size, embed_dims, height, width)
									# with the width and height being successively divided by 2, after every 2 blocks
									for i in range(len(a)):
												self.assertEqual(
												    hidden_states[i].shape		,					torch.Size(
												        [
												            self.model_tester.batch_size,
												            self.model_tester.embed_dims[i // 2],
												            (self.model_tester.image_size // 4) // 2 ** (i // 2),
												            (self.model_tester.image_size // 4) // 2 ** (i // 2),
												        ])		,					)

						SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE						= True
									check_hidden_states_output(a		,					a		,					a)

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

									check_hidden_states_output(a		,					a		,					a)






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						def _config_zero_init(a):
									SCREAMING_SNAKE_CASE						= copy.deepcopy(a)
									for key in configs_no_init.__dict__.keys():
												if "_range" in key or "_std" in key or "initializer_factor" in key or "layer_scale" in key:
															setattr(a		,					a		,					1E-10)
												if isinstance(getattr(a		,					a		,					a)		,					a):
															SCREAMING_SNAKE_CASE						= _config_zero_init(getattr(a		,					a))
															setattr(a		,					a		,					a)
									return configs_no_init

						SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs_for_common()

						SCREAMING_SNAKE_CASE						= _config_zero_init(a)
						for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE						= model_class(config=a)
									for name, param in model.named_parameters():
												if param.requires_grad:
															self.assertIn(
															    ((param.data.mean() * 1E9) / 1E9).round().item()		,					[0.0, 1.0]		,					msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''		,					)







			@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.')
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Dict:
						pass






def 				lowerCamelCase__				():
			SCREAMING_SNAKE_CASE						= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
			return image







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






			@cached_property
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						return ViTImageProcessor.from_pretrained('MBZUAI/swiftformer-xs') if is_vision_available() else None







			@slow
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Union[str, Any]:
						SCREAMING_SNAKE_CASE						= SwiftFormerForImageClassification.from_pretrained('MBZUAI/swiftformer-xs').to(a)

						SCREAMING_SNAKE_CASE						= self.default_image_processor
						SCREAMING_SNAKE_CASE						= prepare_img()
						SCREAMING_SNAKE_CASE						= image_processor(images=a		,					return_tensors='pt').to(a)

						# forward pass
						with torch.no_grad():
									SCREAMING_SNAKE_CASE						= model(**a)

						# verify the logits
						SCREAMING_SNAKE_CASE						= torch.Size((1, 1000))
						self.assertEqual(outputs.logits.shape		,					a)

						SCREAMING_SNAKE_CASE						= torch.tensor([[-2.17_03E00, 2.11_07E00, -2.08_11E00]]).to(a)
						self.assertTrue(torch.allclose(outputs.logits[0, :3]		,					a		,					atol=1E-4))












import copy

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings


a_		:						List[str]             =    R'\n    [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and\n    can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        title_sep (`str`, *optional*, defaults to  `" / "`):\n            Separator inserted between the title and the text of the retrieved document when calling [`RagRetriever`].\n        doc_sep (`str`, *optional*, defaults to  `" // "`):\n            Separator inserted between the text of the retrieved document and the original input when calling\n            [`RagRetriever`].\n        n_docs (`int`, *optional*, defaults to 5):\n            Number of documents to retrieve.\n        max_combined_length (`int`, *optional*, defaults to 300):\n            Max length of contextualized input returned by [`~RagRetriever.__call__`].\n        retrieval_vector_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the document embeddings indexed by [`RagRetriever`].\n        retrieval_batch_size (`int`, *optional*, defaults to 8):\n            Retrieval batch size, defined as the number of queries issues concurrently to the faiss index encapsulated\n            [`RagRetriever`].\n        dataset (`str`, *optional*, defaults to `"wiki_dpr"`):\n            A dataset identifier of the indexed dataset in HuggingFace Datasets (list all available datasets and ids\n            using `datasets.list_datasets()`).\n        dataset_split (`str`, *optional*, defaults to `"train"`)\n            Which split of the `dataset` to load.\n        index_name (`str`, *optional*, defaults to `"compressed"`)\n            The index name of the index associated with the `dataset`. One can choose between `"legacy"`, `"exact"` and\n            `"compressed"`.\n        index_path (`str`, *optional*)\n            The path to the serialized faiss index on disk.\n        passages_path (`str`, *optional*):\n            A path to text passages compatible with the faiss index. Required if using\n            [`~models.rag.retrieval_rag.LegacyIndex`]\n        use_dummy_dataset (`bool`, *optional*, defaults to `False`)\n            Whether to load a "dummy" variant of the dataset specified by `dataset`.\n        label_smoothing (`float`, *optional*, defaults to 0.0):\n            Only relevant if `return_loss` is set to `True`. Controls the `epsilon` parameter value for label smoothing\n            in the loss calculation. If set to 0, no label smoothing is performed.\n        do_marginalize (`bool`, *optional*, defaults to `False`):\n            If `True`, the logits are marginalized over all documents by making use of\n            `torch.nn.functional.log_softmax`.\n        reduce_loss (`bool`, *optional*, defaults to `False`):\n            Whether or not to reduce the NLL loss using the `torch.Tensor.sum` operation.\n        do_deduplication (`bool`, *optional*, defaults to `True`):\n            Whether or not to deduplicate the generations from different context documents for a given input. Has to be\n            set to `False` if used while training with distributed backend.\n        exclude_bos_score (`bool`, *optional*, defaults to `False`):\n            Whether or not to disregard the BOS token when computing the loss.\n        output_retrieved(`bool`, *optional*, defaults to `False`):\n            If set to `True`, `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and\n            `context_attention_mask` are returned. See returned tensors for more detail.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n        forced_eos_token_id (`int`, *optional*):\n            The id of the token to force as the last generated token when `max_length` is reached. Usually set to\n            `eos_token_id`.\n'

@add_start_docstrings(A__					)
class 		_snake_case  (  A__					):
			_lowercase		:    Optional[int]			     =							'''rag'''
			_lowercase		:    Optional[int]			     =							True






			def __init__(       self		,					a=None		,					a=True		,					a=None		,					a=None		,					a=None		,					a=None		,					a=None		,					a=" / "		,					a=" // "		,					a=5		,					a=300		,					a=768		,					a=8		,					a="wiki_dpr"		,					a="train"		,					a="compressed"		,					a=None		,					a=None		,					a=False		,					a=False		,					a=0.0		,					a=True		,					a=False		,					a=False		,					a=False		,					a=True		,					a=None		,					**a		,					)      ->       Optional[Any]:
						super().__init__(
						    bos_token_id=a		,					pad_token_id=a		,					eos_token_id=a		,					decoder_start_token_id=a		,					forced_eos_token_id=a		,					is_encoder_decoder=a		,					prefix=a		,					vocab_size=a		,					**a		,					)
						assert (
						    "question_encoder" in kwargs and "generator" in kwargs
						), "Config has to be initialized with question_encoder and generator config"
						SCREAMING_SNAKE_CASE						= kwargs.pop('question_encoder')
						SCREAMING_SNAKE_CASE						= question_encoder_config.pop('model_type')
						SCREAMING_SNAKE_CASE						= kwargs.pop('generator')
						SCREAMING_SNAKE_CASE						= decoder_config.pop('model_type')

						from ..auto.configuration_auto import AutoConfig

						SCREAMING_SNAKE_CASE						= AutoConfig.for_model(a		,					**a)
						SCREAMING_SNAKE_CASE						= AutoConfig.for_model(a		,					**a)

						SCREAMING_SNAKE_CASE						= reduce_loss
						SCREAMING_SNAKE_CASE						= label_smoothing
						SCREAMING_SNAKE_CASE						= exclude_bos_score
						SCREAMING_SNAKE_CASE						= do_marginalize

						SCREAMING_SNAKE_CASE						= title_sep
						SCREAMING_SNAKE_CASE						= doc_sep
						SCREAMING_SNAKE_CASE						= n_docs
						SCREAMING_SNAKE_CASE						= max_combined_length

						SCREAMING_SNAKE_CASE						= dataset
						SCREAMING_SNAKE_CASE						= dataset_split
						SCREAMING_SNAKE_CASE						= index_name

						SCREAMING_SNAKE_CASE						= retrieval_vector_size
						SCREAMING_SNAKE_CASE						= retrieval_batch_size
						SCREAMING_SNAKE_CASE						= passages_path
						SCREAMING_SNAKE_CASE						= index_path
						SCREAMING_SNAKE_CASE						= use_dummy_dataset

						SCREAMING_SNAKE_CASE						= output_retrieved

						SCREAMING_SNAKE_CASE						= do_deduplication

						SCREAMING_SNAKE_CASE						= use_cache

						if self.forced_eos_token_id is None:
									SCREAMING_SNAKE_CASE						= getattr(self.generator		,					'forced_eos_token_id'		,					a)






			@classmethod
			def    SCREAMING_SNAKE_CASE__    (       cls		,					a		,					a		,					**a)      ->       PretrainedConfig:
						return cls(question_encoder=question_encoder_config.to_dict()		,					generator=generator_config.to_dict()		,					**a)







			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Union[str, Any]:
						SCREAMING_SNAKE_CASE						= copy.deepcopy(self.__dict__)
						SCREAMING_SNAKE_CASE						= self.question_encoder.to_dict()
						SCREAMING_SNAKE_CASE						= self.generator.to_dict()
						SCREAMING_SNAKE_CASE						= self.__class__.model_type
						return output




