
"""simple docstring"""




import json
import os
import shutil
import sys
import tempfile
import unittest
from unittest import TestCase
from unittest.mock import patch

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
					from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
					UpperCAmelCase__				:   List[Any]   							=			None

if is_ray_available():
					import ray  # noqa: E402 # isort:skip
					from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
					UpperCAmelCase__				:   Optional[int]   							=			None
					UpperCAmelCase__				:   str   							=			None
					UpperCAmelCase__				:   List[Any]   							=			None

def lowercase_  (       _snake_case						):
							if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
														SCREAMING_SNAKE_CASE__						:	Dict									=       unittest.skip("""test requires Datasets, Faiss, psutil"""						)(_snake_case						)
							return test_case



@require_distributed_retrieval
class      lowerCAmelCase_						(a__  ):


				"""simple docstring"""



				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       tempfile.mkdtemp()
											SCREAMING_SNAKE_CASE__						:	int									=       8

											# DPR tok
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       [
											    """[UNK]""",
											    """[CLS]""",
											    """[SEP]""",
											    """[PAD]""",
											    """[MASK]""",
											    """want""",
											    """##want""",
											    """##ed""",
											    """wa""",
											    """un""",
											    """runn""",
											    """##ing""",
											    """,""",
											    """low""",
											    """lowest""",
											]
											SCREAMING_SNAKE_CASE__						:	Dict									=       os.path.join(self.tmpdirname       ,	"""dpr_tokenizer"""   )
											os.makedirs(SCREAMING_SNAKE_CASE__       ,	exist_ok=SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	List[str]									=       os.path.join(SCREAMING_SNAKE_CASE__       ,	DPR_VOCAB_FILES_NAMES["""vocab_file"""]   )
											with open(self.vocab_file       ,	"""w"""       ,	encoding="""utf-8"""   ) as vocab_writer:
																		vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]   )   )

											# BART tok
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       [
											    """l""",
											    """o""",
											    """w""",
											    """e""",
											    """r""",
											    """s""",
											    """t""",
											    """i""",
											    """d""",
											    """n""",
											    """\u0120""",
											    """\u0120l""",
											    """\u0120n""",
											    """\u0120lo""",
											    """\u0120low""",
											    """er""",
											    """\u0120lowest""",
											    """\u0120newer""",
											    """\u0120wider""",
											    """<unk>""",
											]
											SCREAMING_SNAKE_CASE__						:	int									=       dict(zip(SCREAMING_SNAKE_CASE__       ,	range(len(SCREAMING_SNAKE_CASE__   )   )   )   )
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       ["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
											SCREAMING_SNAKE_CASE__						:	int									=       {"""unk_token""": """<unk>"""}

											SCREAMING_SNAKE_CASE__						:	List[str]									=       os.path.join(self.tmpdirname       ,	"""bart_tokenizer"""   )
											os.makedirs(SCREAMING_SNAKE_CASE__       ,	exist_ok=SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Any									=       os.path.join(SCREAMING_SNAKE_CASE__       ,	BART_VOCAB_FILES_NAMES["""vocab_file"""]   )
											SCREAMING_SNAKE_CASE__						:	str									=       os.path.join(SCREAMING_SNAKE_CASE__       ,	BART_VOCAB_FILES_NAMES["""merges_file"""]   )
											with open(self.vocab_file       ,	"""w"""       ,	encoding="""utf-8"""   ) as fp:
																		fp.write(json.dumps(SCREAMING_SNAKE_CASE__   ) + """\n"""   )
											with open(self.merges_file       ,	"""w"""       ,	encoding="""utf-8"""   ) as fp:
																		fp.write("""\n""".join(SCREAMING_SNAKE_CASE__   )   )



				def      __magic_name__   (self   ):



											"""simple docstring"""




											return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,	"""dpr_tokenizer"""   )   )



				def      __magic_name__   (self   ):



											"""simple docstring"""




											return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,	"""bart_tokenizer"""   )   )



				def      __magic_name__   (self   ):



											"""simple docstring"""




											shutil.rmtree(self.tmpdirname   )



				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Tuple									=       Dataset.from_dict(
											    {
											        """id""": ["""0""", """1"""],
											        """text""": ["""foo""", """bar"""],
											        """title""": ["""Foo""", """Bar"""],
											        """embeddings""": [np.ones(self.retrieval_vector_size   ), 2 * np.ones(self.retrieval_vector_size   )],
											    }   )
											dataset.add_faiss_index("""embeddings"""       ,	string_factory="""Flat"""       ,	metric_type=faiss.METRIC_INNER_PRODUCT   )
											return dataset



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__=1_23_45   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	int									=       self.get_dummy_dataset()
											SCREAMING_SNAKE_CASE__						:	Tuple									=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       ,	question_encoder=DPRConfig().to_dict()       ,	generator=BartConfig().to_dict()       ,	)
											with patch("""transformers.models.rag.retrieval_rag.load_dataset"""   ) as mock_load_dataset:
																		SCREAMING_SNAKE_CASE__						:	Any									=       dataset
																		SCREAMING_SNAKE_CASE__						:	Tuple									=       RagPyTorchDistributedRetriever(
																		    SCREAMING_SNAKE_CASE__       ,	question_encoder_tokenizer=self.get_dpr_tokenizer()       ,	generator_tokenizer=self.get_bart_tokenizer()       ,	)
																		if init_retrieval:
																									retriever.init_retrieval(SCREAMING_SNAKE_CASE__   )
											return retriever



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__   ):



											"""simple docstring"""




											ray.init(local_mode=SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Any									=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       ,	question_encoder=DPRConfig().to_dict()       ,	generator=BartConfig().to_dict()       ,	)
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       ray.remote(SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       [remote_cls.remote() for _ in range(1   )]
											with patch("""transformers.models.rag.retrieval_rag.load_dataset"""   ) as mock_load_dataset:
																		SCREAMING_SNAKE_CASE__						:	List[str]									=       self.get_dummy_dataset()
																		SCREAMING_SNAKE_CASE__						:	Optional[int]									=       RagRayDistributedRetriever(
																		    SCREAMING_SNAKE_CASE__       ,	question_encoder_tokenizer=self.get_dpr_tokenizer()       ,	generator_tokenizer=self.get_bart_tokenizer()       ,	retrieval_workers=SCREAMING_SNAKE_CASE__       ,	)
																		if init_retrieval:
																									retriever.init_retrieval()
											return retriever



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__=1_23_45   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	List[str]									=       self.get_dummy_dataset()
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       ,	question_encoder=DPRConfig().to_dict()       ,	generator=BartConfig().to_dict()       ,	index_name="""custom"""       ,	)
											if from_disk:
																		SCREAMING_SNAKE_CASE__						:	List[Any]									=       os.path.join(self.tmpdirname       ,	"""dataset"""   )
																		SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       os.path.join(self.tmpdirname       ,	"""index.faiss"""   )
																		dataset.get_index("""embeddings"""   ).save(os.path.join(self.tmpdirname       ,	"""index.faiss"""   )   )
																		dataset.drop_index("""embeddings"""   )
																		dataset.save_to_disk(os.path.join(self.tmpdirname       ,	"""dataset"""   )   )
																		del dataset
																		SCREAMING_SNAKE_CASE__						:	Optional[int]									=       RagPyTorchDistributedRetriever(
																		    SCREAMING_SNAKE_CASE__       ,	question_encoder_tokenizer=self.get_dpr_tokenizer()       ,	generator_tokenizer=self.get_bart_tokenizer()       ,	)
											else:
																		SCREAMING_SNAKE_CASE__						:	str									=       RagPyTorchDistributedRetriever(
																		    SCREAMING_SNAKE_CASE__       ,	question_encoder_tokenizer=self.get_dpr_tokenizer()       ,	generator_tokenizer=self.get_bart_tokenizer()       ,	index=CustomHFIndex(config.retrieval_vector_size       ,	SCREAMING_SNAKE_CASE__   )       ,	)
											if init_retrieval:
																		retriever.init_retrieval(SCREAMING_SNAKE_CASE__   )
											return retriever



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ):



											"""simple docstring"""




											ray.init(local_mode=SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       self.get_dummy_dataset()
											SCREAMING_SNAKE_CASE__						:	Dict									=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       ,	question_encoder=DPRConfig().to_dict()       ,	generator=BartConfig().to_dict()       ,	index_name="""custom"""       ,	)
											SCREAMING_SNAKE_CASE__						:	Any									=       ray.remote(SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       [remote_cls.remote() for _ in range(1   )]
											if from_disk:
																		SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       os.path.join(self.tmpdirname       ,	"""dataset"""   )
																		SCREAMING_SNAKE_CASE__						:	Dict									=       os.path.join(self.tmpdirname       ,	"""index.faiss"""   )
																		dataset.get_index("""embeddings"""   ).save(os.path.join(self.tmpdirname       ,	"""index.faiss"""   )   )
																		dataset.drop_index("""embeddings"""   )
																		dataset.save_to_disk(os.path.join(self.tmpdirname       ,	"""dataset"""   )   )
																		del dataset
																		SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       RagRayDistributedRetriever(
																		    SCREAMING_SNAKE_CASE__       ,	question_encoder_tokenizer=self.get_dpr_tokenizer()       ,	generator_tokenizer=self.get_bart_tokenizer()       ,	retrieval_workers=SCREAMING_SNAKE_CASE__       ,	index=CustomHFIndex.load_from_disk(
																		        vector_size=config.retrieval_vector_size       ,	dataset_path=config.passages_path       ,	index_path=config.index_path       ,	)       ,	)
											else:
																		SCREAMING_SNAKE_CASE__						:	int									=       RagRayDistributedRetriever(
																		    SCREAMING_SNAKE_CASE__       ,	question_encoder_tokenizer=self.get_dpr_tokenizer()       ,	generator_tokenizer=self.get_bart_tokenizer()       ,	retrieval_workers=SCREAMING_SNAKE_CASE__       ,	index=CustomHFIndex(config.retrieval_vector_size       ,	SCREAMING_SNAKE_CASE__   )       ,	)
											if init_retrieval:
																		retriever.init_retrieval()
											return retriever



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Any									=       retriever.retrieve(SCREAMING_SNAKE_CASE__       ,	n_docs=SCREAMING_SNAKE_CASE__   )
											self.assertEqual(retrieved_doc_embeds.shape       ,	(2, n_docs, self.retrieval_vector_size)   )
											self.assertEqual(len(SCREAMING_SNAKE_CASE__   )       ,	2   )
											self.assertEqual(sorted(doc_dicts[0]   )       ,	["""embeddings""", """id""", """text""", """title"""]   )
											self.assertEqual(len(doc_dicts[0]["""id"""]   )       ,	SCREAMING_SNAKE_CASE__   )
											self.assertEqual(doc_dicts[0]["""id"""][0]       ,	"""1"""   )  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]["""id"""][0]       ,	"""0"""   )  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist()       ,	[[1], [0]]   )



				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Tuple									=       1
											SCREAMING_SNAKE_CASE__						:	str									=       np.array(
											    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,	dtype=np.floataa   )

											self.distributed_retriever_check(
											    self.get_dummy_pytorch_distributed_retriever(init_retrieval=SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Dict									=       1
											SCREAMING_SNAKE_CASE__						:	List[str]									=       np.array(
											    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,	dtype=np.floataa   )

											self.distributed_retriever_check(
											    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=SCREAMING_SNAKE_CASE__       ,	from_disk=SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	)



				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	List[Any]									=       1
											SCREAMING_SNAKE_CASE__						:	str									=       np.array(
											    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,	dtype=np.floataa   )

											self.distributed_retriever_check(
											    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=SCREAMING_SNAKE_CASE__       ,	from_disk=SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	)



				@require_ray
				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Any									=       1
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       np.array(
											    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,	dtype=np.floataa   )

											self.distributed_retriever_check(
											    self.get_dummy_ray_distributed_retriever(init_retrieval=SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )
											ray.shutdown()



				@require_ray
				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       1
											SCREAMING_SNAKE_CASE__						:	Any									=       np.array(
											    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,	dtype=np.floataa   )
											with self.assertRaises(SCREAMING_SNAKE_CASE__   ):
																		self.distributed_retriever_check(
																		    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=SCREAMING_SNAKE_CASE__       ,	from_disk=SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	)
											ray.shutdown()



				@require_ray
				def      __magic_name__   (self   ):



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       1
											SCREAMING_SNAKE_CASE__						:	int									=       np.array(
											    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,	dtype=np.floataa   )

											self.distributed_retriever_check(
											    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=SCREAMING_SNAKE_CASE__       ,	from_disk=SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )
											ray.shutdown()






"""simple docstring"""




import collections
import inspect
import unittest

from transformers import DonutSwinConfig
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import is_torch_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 DonutSwinModel
					from transformers.models.donut.modeling_donut_swin import DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST



class      lowerCAmelCase_						:


				"""simple docstring"""



				def __init__(self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__=13       ,	SCREAMING_SNAKE_CASE__=32       ,	SCREAMING_SNAKE_CASE__=2       ,	SCREAMING_SNAKE_CASE__=3       ,	SCREAMING_SNAKE_CASE__=16       ,	SCREAMING_SNAKE_CASE__=[1, 2, 1]       ,	SCREAMING_SNAKE_CASE__=[2, 2, 4]       ,	SCREAMING_SNAKE_CASE__=2       ,	SCREAMING_SNAKE_CASE__=2.0       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=0.0       ,	SCREAMING_SNAKE_CASE__=0.0       ,	SCREAMING_SNAKE_CASE__=0.1       ,	SCREAMING_SNAKE_CASE__="gelu"       ,	SCREAMING_SNAKE_CASE__=False       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=0.02       ,	SCREAMING_SNAKE_CASE__=1E-5       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=None       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=10       ,	SCREAMING_SNAKE_CASE__=8       ,	) ->		List[str]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	List[str]									=       parent
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       batch_size
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       image_size
											SCREAMING_SNAKE_CASE__						:	List[str]									=       patch_size
											SCREAMING_SNAKE_CASE__						:	List[str]									=       num_channels
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       embed_dim
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       depths
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       num_heads
											SCREAMING_SNAKE_CASE__						:	str									=       window_size
											SCREAMING_SNAKE_CASE__						:	List[str]									=       mlp_ratio
											SCREAMING_SNAKE_CASE__						:	Tuple									=       qkv_bias
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       hidden_dropout_prob
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE__						:	List[str]									=       drop_path_rate
											SCREAMING_SNAKE_CASE__						:	str									=       hidden_act
											SCREAMING_SNAKE_CASE__						:	Any									=       use_absolute_embeddings
											SCREAMING_SNAKE_CASE__						:	Any									=       patch_norm
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       layer_norm_eps
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       initializer_range
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       is_training
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       scope
											SCREAMING_SNAKE_CASE__						:	List[str]									=       use_labels
											SCREAMING_SNAKE_CASE__						:	str									=       type_sequence_label_size
											SCREAMING_SNAKE_CASE__						:	str									=       encoder_stride



				def      __magic_name__   (self   ) ->		Union[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

											SCREAMING_SNAKE_CASE__						:	Any									=       None
											if self.use_labels:
																		SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       ids_tensor([self.batch_size]       ,	self.type_sequence_label_size   )

											SCREAMING_SNAKE_CASE__						:	List[str]									=       self.get_config()

											return config, pixel_values, labels



				def      __magic_name__   (self   ) ->		Optional[Any]:



											"""simple docstring"""




											return DonutSwinConfig(
											    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      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		int:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       DonutSwinModel(config=SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       model(SCREAMING_SNAKE_CASE__   )

											SCREAMING_SNAKE_CASE__						:	Any									=       ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths   ) - 1))
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       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      __magic_name__   (self   ) ->		Any:



											"""simple docstring"""




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






@require_torch
class      lowerCAmelCase_						(a__			, a__			, unittest.TestCase  ):


				"""simple docstring"""



				__UpperCamelCase				:						str 	=       (DonutSwinModel,) if is_torch_available() else ()
				__UpperCamelCase				:						str 	=       {'''feature-extraction''': DonutSwinModel} if is_torch_available() else {}
				__UpperCamelCase				:						List[str] 	=       True

				__UpperCamelCase				:						str 	=       False
				__UpperCamelCase				:						List[Any] 	=       False
				__UpperCamelCase				:						List[Any] 	=       False



				def      __magic_name__   (self   ) ->		Tuple:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Any									=       DonutSwinModelTester(self   )
											SCREAMING_SNAKE_CASE__						:	Dict									=       ConfigTester(self       ,	config_class=SCREAMING_SNAKE_CASE__       ,	embed_dim=37   )



				def      __magic_name__   (self   ) ->		int:



											"""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      __magic_name__   (self   ) ->		Optional[Any]:



											"""simple docstring"""




											return



				def      __magic_name__   (self   ) ->		Tuple:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		List[str]:



											"""simple docstring"""




											pass



				def      __magic_name__   (self   ) ->		List[str]:



											"""simple docstring"""




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

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



				def      __magic_name__   (self   ) ->		Any:



											"""simple docstring"""




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

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

																		SCREAMING_SNAKE_CASE__						:	str									=       ["""pixel_values"""]
																		self.assertListEqual(arg_names[:1]       ,	SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		Optional[Any]:



											"""simple docstring"""




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

											for model_class in self.all_model_classes:
																		SCREAMING_SNAKE_CASE__						:	List[Any]									=       True
																		SCREAMING_SNAKE_CASE__						:	Any									=       False
																		SCREAMING_SNAKE_CASE__						:	str									=       True
																		SCREAMING_SNAKE_CASE__						:	Optional[int]									=       model_class(SCREAMING_SNAKE_CASE__   )
																		model.to(SCREAMING_SNAKE_CASE__   )
																		model.eval()
																		with torch.no_grad():
																									SCREAMING_SNAKE_CASE__						:	List[Any]									=       model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )   )
																		SCREAMING_SNAKE_CASE__						:	str									=       outputs.attentions
																		SCREAMING_SNAKE_CASE__						:	List[Any]									=       len(self.model_tester.depths   )
																		self.assertEqual(len(SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__   )

																		# check that output_attentions also work using config
																		del inputs_dict["output_attentions"]
																		SCREAMING_SNAKE_CASE__						:	Optional[int]									=       True
																		SCREAMING_SNAKE_CASE__						:	int									=       config.window_size**2
																		SCREAMING_SNAKE_CASE__						:	int									=       model_class(SCREAMING_SNAKE_CASE__   )
																		model.to(SCREAMING_SNAKE_CASE__   )
																		model.eval()
																		with torch.no_grad():
																									SCREAMING_SNAKE_CASE__						:	Tuple									=       model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )   )
																		SCREAMING_SNAKE_CASE__						:	Tuple									=       outputs.attentions
																		self.assertEqual(len(SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__   )

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

																		# Check attention is always last and order is fine
																		SCREAMING_SNAKE_CASE__						:	List[str]									=       True
																		SCREAMING_SNAKE_CASE__						:	str									=       True
																		SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       model_class(SCREAMING_SNAKE_CASE__   )
																		model.to(SCREAMING_SNAKE_CASE__   )
																		model.eval()
																		with torch.no_grad():
																									SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )   )

																		if hasattr(self.model_tester       ,	"""num_hidden_states_types"""   ):
																									SCREAMING_SNAKE_CASE__						:	str									=       self.model_tester.num_hidden_states_types
																		else:
																									# also another +1 for reshaped_hidden_states
																									SCREAMING_SNAKE_CASE__						:	List[Any]									=       2
																		self.assertEqual(out_len + added_hidden_states       ,	len(SCREAMING_SNAKE_CASE__   )   )

																		SCREAMING_SNAKE_CASE__						:	List[Any]									=       outputs.attentions

																		self.assertEqual(len(SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__   )

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



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	List[Any]									=       model_class(SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()

											with torch.no_grad():
																		SCREAMING_SNAKE_CASE__						:	str									=       model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )   )

											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       outputs.hidden_states

											SCREAMING_SNAKE_CASE__						:	List[str]									=       getattr(
											    self.model_tester       ,	"""expected_num_hidden_layers"""       ,	len(self.model_tester.depths   ) + 1   )
											self.assertEqual(len(SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__   )

											# DonutSwin has a different seq_length
											SCREAMING_SNAKE_CASE__						:	Any									=       (
											    config.patch_size
											    if isinstance(config.patch_size       ,	collections.abc.Iterable   )
											    else (config.patch_size, config.patch_size)
											)

											SCREAMING_SNAKE_CASE__						:	int									=       (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]       ,	)

											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       outputs.reshaped_hidden_states
											self.assertEqual(len(SCREAMING_SNAKE_CASE__   )       ,	SCREAMING_SNAKE_CASE__   )

											SCREAMING_SNAKE_CASE__				,				SCREAMING_SNAKE_CASE__				,				SCREAMING_SNAKE_CASE__				,				SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       reshaped_hidden_states[0].shape
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       (
											    reshaped_hidden_states[0].view(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	height * width   ).permute(0       ,	2       ,	1   )
											)
											self.assertListEqual(
											    list(reshaped_hidden_states.shape[-2:]   )       ,	[num_patches, self.model_tester.embed_dim]       ,	)



				def      __magic_name__   (self   ) ->		Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__				,				SCREAMING_SNAKE_CASE__						:	Any									=       self.model_tester.prepare_config_and_inputs_for_common()

											SCREAMING_SNAKE_CASE__						:	Tuple									=       (
											    self.model_tester.image_size
											    if isinstance(self.model_tester.image_size       ,	collections.abc.Iterable   )
											    else (self.model_tester.image_size, self.model_tester.image_size)
											)

											for model_class in self.all_model_classes:
																		SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       True
																		self.check_hidden_states_output(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )

																		# check that output_hidden_states also work using config
																		del inputs_dict["output_hidden_states"]
																		SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       True

																		self.check_hidden_states_output(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__				,				SCREAMING_SNAKE_CASE__						:	Dict									=       self.model_tester.prepare_config_and_inputs_for_common()
											SCREAMING_SNAKE_CASE__						:	Tuple									=       3

											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       (
											    self.model_tester.image_size
											    if isinstance(self.model_tester.image_size       ,	collections.abc.Iterable   )
											    else (self.model_tester.image_size, self.model_tester.image_size)
											)
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       (
											    config.patch_size
											    if isinstance(config.patch_size       ,	collections.abc.Iterable   )
											    else (config.patch_size, config.patch_size)
											)

											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

											for model_class in self.all_model_classes:
																		SCREAMING_SNAKE_CASE__						:	Any									=       True
																		self.check_hidden_states_output(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	(padded_height, padded_width)   )

																		# check that output_hidden_states also work using config
																		del inputs_dict["output_hidden_states"]
																		SCREAMING_SNAKE_CASE__						:	List[str]									=       True
																		self.check_hidden_states_output(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	(padded_height, padded_width)   )



				@slow
				def      __magic_name__   (self   ) ->		Union[str, Any]:



											"""simple docstring"""




											for model_name in DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		SCREAMING_SNAKE_CASE__						:	Optional[int]									=       DonutSwinModel.from_pretrained(SCREAMING_SNAKE_CASE__   )
																		self.assertIsNotNone(SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__				,				SCREAMING_SNAKE_CASE__						:	str									=       self.model_tester.prepare_config_and_inputs_for_common()

											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       _config_zero_init(SCREAMING_SNAKE_CASE__   )
											for model_class in self.all_model_classes:
																		SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       model_class(config=SCREAMING_SNAKE_CASE__   )
																		for name, param in model.named_parameters():
																									if "embeddings" not in name and 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'''       ,	)




