





'''simple docstring'''






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

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
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_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
		import faiss







@require_faiss
class 			lowerCamelCase					(						UpperCamelCase__  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   Tuple  					=		tempfile.mkdtemp()
											_snake_case:   List[Any]  					=		8

											# DPR tok
											_snake_case:   Optional[Any]  					=		[
											    '[UNK]',
											    '[CLS]',
											    '[SEP]',
											    '[PAD]',
											    '[MASK]',
											    'want',
											    '##want',
											    '##ed',
											    'wa',
											    'un',
											    'runn',
											    '##ing',
											    ',',
											    'low',
											    'lowest',
											]
											_snake_case:   int  					=		os.path.join(self.tmpdirname       ,     'dpr_tokenizer'       )
											os.makedirs(__A       ,     exist_ok=__A       )
											_snake_case:   Optional[Any]  					=		os.path.join(__A       ,     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
											_snake_case:   str  					=		[
											    'l',
											    'o',
											    'w',
											    'e',
											    'r',
											    's',
											    't',
											    'i',
											    'd',
											    'n',
											    '\u0120',
											    '\u0120l',
											    '\u0120n',
											    '\u0120lo',
											    '\u0120low',
											    'er',
											    '\u0120lowest',
											    '\u0120newer',
											    '\u0120wider',
											    '<unk>',
											]
											_snake_case:   List[str]  					=		dict(zip(__A       ,     range(len(__A       )       )       )       )
											_snake_case:   str  					=		['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
											_snake_case:   List[str]  					=		{'unk_token': '<unk>'}

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


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





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


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,     'dpr_tokenizer'       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





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


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											shutil.rmtree(self.tmpdirname       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_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    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.get_dummy_dataset()
											_snake_case:   Optional[int]  					=		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:
																	_snake_case:   List[str]  					=		dataset
																	_snake_case:   int  					=		RagRetriever(
																	    __A       ,     question_encoder_tokenizer=self.get_dpr_tokenizer()       ,     generator_tokenizer=self.get_bart_tokenizer()       ,     )
											return retriever


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ,     __snake_case	:  bool       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.get_dummy_dataset()
											_snake_case:   str  					=		RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       ,     question_encoder=DPRConfig().to_dict()       ,     generator=BartConfig().to_dict()       ,     index_name='custom'       ,     )
											if from_disk:
																	_snake_case:   str  					=		os.path.join(self.tmpdirname       ,     'dataset'       )
																	_snake_case:   Optional[int]  					=		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
																	_snake_case:   Any  					=		RagRetriever(
																	    __A       ,     question_encoder_tokenizer=self.get_dpr_tokenizer()       ,     generator_tokenizer=self.get_bart_tokenizer()       ,     )
											else:
																	_snake_case:   Dict  					=		RagRetriever(
																	    __A       ,     question_encoder_tokenizer=self.get_dpr_tokenizer()       ,     generator_tokenizer=self.get_bart_tokenizer()       ,     index=CustomHFIndex(config.retrieval_vector_size       ,     __A       )       ,     )
											return retriever


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





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

											_snake_case:   Optional[int]  					=		os.path.join(self.tmpdirname       ,     'hf_bert_base.hnswSQ8_correct_phi_128.c_index'       )
											dataset.save_faiss_index('embeddings'       ,     index_file_name + '.index.dpr'       )
											pickle.dump(dataset['id']       ,     open(index_file_name + '.index_meta.dpr'       ,     'wb'       )       )

											_snake_case:   List[Any]  					=		os.path.join(self.tmpdirname       ,     'psgs_w100.tsv.pkl'       )
											_snake_case:   List[Any]  					=		{sample['id']: [sample['text'], sample['title']] for sample in dataset}
											pickle.dump(__A       ,     open(__A       ,     'wb'       )       )

											_snake_case:   Any  					=		RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       ,     question_encoder=DPRConfig().to_dict()       ,     generator=BartConfig().to_dict()       ,     index_name='legacy'       ,     index_path=self.tmpdirname       ,     )
											_snake_case:   Any  					=		RagRetriever(
											    __A       ,     question_encoder_tokenizer=self.get_dpr_tokenizer()       ,     generator_tokenizer=self.get_bart_tokenizer()       )
											return retriever


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											_snake_case:   Tuple  					=		1
											_snake_case:   Any  					=		self.get_dummy_canonical_hf_index_retriever()
											_snake_case:   List[Any]  					=		np.array(
											    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
											_snake_case					,						_snake_case					,						_snake_case:   List[str]  					=		retriever.retrieve(__A       ,     n_docs=__A       )
											self.assertEqual(retrieved_doc_embeds.shape       ,     (2, n_docs, self.retrieval_vector_size)       )
											self.assertEqual(len(__A       )       ,     2       )
											self.assertEqual(sorted(doc_dicts[0]       )       ,     ['embeddings', 'id', 'text', 'title']       )
											self.assertEqual(len(doc_dicts[0]['id']       )       ,     __A       )
											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    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		self.get_dummy_canonical_hf_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																	with patch('transformers.models.rag.retrieval_rag.load_dataset'       ) as mock_load_dataset:
																							_snake_case:   Optional[int]  					=		self.get_dummy_dataset()
																							retriever.save_pretrained(__A       )
																							_snake_case:   List[Any]  					=		RagRetriever.from_pretrained(__A       )
																							self.assertIsInstance(__A       ,     __A       )
																							_snake_case:   str  					=		np.array(
																							    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
																							_snake_case:   List[str]  					=		retriever.retrieve(__A       ,     n_docs=1       )
																							self.assertTrue(out is not None       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   List[Any]  					=		1
											_snake_case:   List[Any]  					=		self.get_dummy_custom_hf_index_retriever(from_disk=__A       )
											_snake_case:   Any  					=		np.array(
											    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
											_snake_case					,						_snake_case					,						_snake_case:   Tuple  					=		retriever.retrieve(__A       ,     n_docs=__A       )
											self.assertEqual(retrieved_doc_embeds.shape       ,     (2, n_docs, self.retrieval_vector_size)       )
											self.assertEqual(len(__A       )       ,     2       )
											self.assertEqual(sorted(doc_dicts[0]       )       ,     ['embeddings', 'id', 'text', 'title']       )
											self.assertEqual(len(doc_dicts[0]['id']       )       ,     __A       )
											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    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		self.get_dummy_custom_hf_index_retriever(from_disk=__A       )
											with tempfile.TemporaryDirectory() as tmp_dirname:
																	retriever.save_pretrained(__A       )
																	_snake_case:   Union[str, Any]  					=		RagRetriever.from_pretrained(__A       )
																	self.assertIsInstance(__A       ,     __A       )
																	_snake_case:   Optional[int]  					=		np.array(
																	    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
																	_snake_case:   Any  					=		retriever.retrieve(__A       ,     n_docs=1       )
																	self.assertTrue(out is not None       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		1
											_snake_case:   Any  					=		self.get_dummy_custom_hf_index_retriever(from_disk=__A       )
											_snake_case:   Union[str, Any]  					=		np.array(
											    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
											_snake_case					,						_snake_case					,						_snake_case:   Optional[Any]  					=		retriever.retrieve(__A       ,     n_docs=__A       )
											self.assertEqual(retrieved_doc_embeds.shape       ,     (2, n_docs, self.retrieval_vector_size)       )
											self.assertEqual(len(__A       )       ,     2       )
											self.assertEqual(sorted(doc_dicts[0]       )       ,     ['embeddings', 'id', 'text', 'title']       )
											self.assertEqual(len(doc_dicts[0]['id']       )       ,     __A       )
											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    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   List[str]  					=		self.get_dummy_custom_hf_index_retriever(from_disk=__A       )
											with tempfile.TemporaryDirectory() as tmp_dirname:
																	retriever.save_pretrained(__A       )
																	_snake_case:   Union[str, Any]  					=		RagRetriever.from_pretrained(__A       )
																	self.assertIsInstance(__A       ,     __A       )
																	_snake_case:   str  					=		np.array(
																	    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
																	_snake_case:   Optional[int]  					=		retriever.retrieve(__A       ,     n_docs=1       )
																	self.assertTrue(out is not None       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											_snake_case:   List[Any]  					=		1
											_snake_case:   Union[str, Any]  					=		self.get_dummy_legacy_index_retriever()
											_snake_case:   Optional[Any]  					=		np.array(
											    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
											_snake_case					,						_snake_case					,						_snake_case:   Any  					=		retriever.retrieve(__A       ,     n_docs=__A       )
											self.assertEqual(retrieved_doc_embeds.shape       ,     (2, n_docs, self.retrieval_vector_size)       )
											self.assertEqual(len(__A       )       ,     2       )
											self.assertEqual(sorted(doc_dicts[0]       )       ,     ['text', 'title']       )
											self.assertEqual(len(doc_dicts[0]['text']       )       ,     __A       )
											self.assertEqual(doc_dicts[0]['text'][0]       ,     'bar'       )  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['text'][0]       ,     'foo'       )  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist()       ,     [[1], [0]]       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   List[str]  					=		self.get_dummy_legacy_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																	retriever.save_pretrained(__A       )
																	_snake_case:   Tuple  					=		RagRetriever.from_pretrained(__A       )
																	self.assertIsInstance(__A       ,     __A       )
																	_snake_case:   str  					=		np.array(
																	    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
																	_snake_case:   Tuple  					=		retriever.retrieve(__A       ,     n_docs=1       )
																	self.assertTrue(out is not None       )


					@require_torch
					@require_tokenizers
					@require_sentencepiece
					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											import torch

											_snake_case:   List[Any]  					=		1
											_snake_case:   Optional[int]  					=		self.get_dummy_canonical_hf_index_retriever()
											_snake_case:   str  					=		[[5, 7], [10, 11]]
											_snake_case:   List[str]  					=		np.array(
											    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
											_snake_case:   int  					=		retriever(__A       ,     __A       ,     prefix=retriever.config.generator.prefix       ,     n_docs=__A       )
											_snake_case					,						_snake_case					,						_snake_case:   int  					=		(
											    out['context_input_ids'],
											    out['context_attention_mask'],
											    out['retrieved_doc_embeds'],
											)
											self.assertEqual(retrieved_doc_embeds.shape       ,     (2, n_docs, self.retrieval_vector_size)       )
											self.assertIsInstance(__A       ,     __A       )
											self.assertIsInstance(__A       ,     __A       )
											self.assertIsInstance(__A       ,     np.ndarray       )

											_snake_case:   List[str]  					=		retriever(
											    __A       ,     __A       ,     prefix=retriever.config.generator.prefix       ,     n_docs=__A       ,     return_tensors='pt'       ,     )
											_snake_case					,						_snake_case					,						_snake_case					,						_snake_case:   Union[str, Any]  					=		(  # noqa: F841
											    out['context_input_ids'],
											    out['context_attention_mask'],
											    out['retrieved_doc_embeds'],
											    out['doc_ids'],
											)
											self.assertEqual(retrieved_doc_embeds.shape       ,     (2, n_docs, self.retrieval_vector_size)       )
											self.assertIsInstance(__A       ,     torch.Tensor       )
											self.assertIsInstance(__A       ,     torch.Tensor       )
											self.assertIsInstance(__A       ,     torch.Tensor       )




					@require_torch
					@require_tokenizers
					@require_sentencepiece
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		self.get_dpr_ctx_encoder_tokenizer()
											_snake_case:   Dict  					=		1
											_snake_case:   str  					=		self.get_dummy_custom_hf_index_retriever(from_disk=__A       )
											retriever.set_ctx_encoder_tokenizer(__A       )

											_snake_case:   List[str]  					=		[[5, 7], [10, 11]]
											_snake_case:   Union[str, Any]  					=		np.array(
											    [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]       ,     dtype=np.floataa       )
											_snake_case:   str  					=		retriever(__A       ,     __A       ,     prefix=retriever.config.generator.prefix       ,     n_docs=__A       )

											self.assertEqual(
											    len(__A       )       ,     6       )  # check whether the retriever output consist of 6 attributes including tokenized docs
											self.assertEqual(
											    all(k in out for k in ('tokenized_doc_ids', 'tokenized_doc_attention_mask')       )       ,     __A       )  # check for doc token related keys in dictionary.











'''simple docstring'''






import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
		import torch
		from torch import nn

		from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
		from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import ViTImageProcessor







class 			lowerCamelCase					:


					def __init__(					self	:  Union[str, Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  str=13       ,     __snake_case	:  Union[str, Any]=30       ,     __snake_case	:  Union[str, Any]=2       ,     __snake_case	:  Dict=3       ,     __snake_case	:  Optional[Any]=True       ,     __snake_case	:  Optional[int]=True       ,     __snake_case	:  Optional[int]=32       ,     __snake_case	:  Optional[int]=5       ,     __snake_case	:  Any=4       ,     __snake_case	:  int=37       ,     __snake_case	:  int="gelu"       ,     __snake_case	:  Union[str, Any]=0.1       ,     __snake_case	:  Optional[int]=0.1       ,     __snake_case	:  int=10       ,     __snake_case	:  Any=0.02       ,     __snake_case	:  List[str]=None       ,     __snake_case	:  Tuple=2       ,     ):





											'''simple docstring'''





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

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


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											_snake_case:   List[str]  					=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

											_snake_case:   List[str]  					=		None
											if self.use_labels:
																	_snake_case:   Any  					=		ids_tensor([self.batch_size]       ,     self.type_sequence_label_size       )

											_snake_case:   Union[str, Any]  					=		self.get_config()

											return config, pixel_values, labels


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											return ViTConfig(
											    image_size=self.image_size       ,     patch_size=self.patch_size       ,     num_channels=self.num_channels       ,     hidden_size=self.hidden_size       ,     num_hidden_layers=self.num_hidden_layers       ,     num_attention_heads=self.num_attention_heads       ,     intermediate_size=self.intermediate_size       ,     hidden_act=self.hidden_act       ,     hidden_dropout_prob=self.hidden_dropout_prob       ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,     is_decoder=__snake_case       ,     initializer_range=self.initializer_range       ,     encoder_stride=self.encoder_stride       ,     )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ,     __snake_case	:  int       ,     __snake_case	:  Optional[int]       ,     __snake_case	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		ViTModel(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											_snake_case:   Tuple  					=		model(__snake_case       )
											self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  int       ):





											'''simple docstring'''





											_snake_case:   int  					=		ViTForMaskedImageModeling(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											_snake_case:   Dict  					=		model(__snake_case       )
											self.parent.assertEqual(
											    result.reconstruction.shape       ,     (self.batch_size, self.num_channels, self.image_size, self.image_size)       )

											# test greyscale images
											_snake_case:   List[str]  					=		1
											_snake_case:   Tuple  					=		ViTForMaskedImageModeling(__snake_case       )
											model.to(__snake_case       )
											model.eval()

											_snake_case:   List[str]  					=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.parent.assertEqual(result.reconstruction.shape       ,     (self.batch_size, 1, self.image_size, self.image_size)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Optional[int]       ,     __snake_case	:  Dict       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		self.type_sequence_label_size
											_snake_case:   Union[str, Any]  					=		ViTForImageClassification(__snake_case       )
											model.to(__snake_case       )
											model.eval()
											_snake_case:   List[Any]  					=		model(__snake_case       ,     labels=__snake_case       )
											self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.type_sequence_label_size)       )

											# test greyscale images
											_snake_case:   Tuple  					=		1
											_snake_case:   Optional[int]  					=		ViTForImageClassification(__snake_case       )
											model.to(__snake_case       )
											model.eval()

											_snake_case:   List[Any]  					=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
											_snake_case:   Any  					=		model(__snake_case       )
											self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.type_sequence_label_size)       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.prepare_config_and_inputs()
											(
											    (
											    _snake_case
											)					,						(
											    _snake_case
											)					,						(
											    _snake_case
											)					,						
											):   int  					=		config_and_inputs
											_snake_case:   Dict  					=		{'pixel_values': pixel_values}
											return config, inputs_dict







@require_torch
class 			lowerCamelCase					(						__UpperCAmelCase     ,							__UpperCAmelCase     ,							unittest.TestCase  ):
					_SCREAMING_SNAKE_CASE							   =    (
					    (
					        ViTModel,
					        ViTForImageClassification,
					        ViTForMaskedImageModeling,
					    )
					    if is_torch_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE							   =    (
					    {"feature-extraction": ViTModel, "image-classification": ViTForImageClassification}
					    if is_torch_available()
					    else {}
					)
					_SCREAMING_SNAKE_CASE							   =    True

					_SCREAMING_SNAKE_CASE							   =    False
					_SCREAMING_SNAKE_CASE							   =    False
					_SCREAMING_SNAKE_CASE							   =    False


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		ViTModelTester(self       )
											_snake_case:   Union[str, Any]  					=		ConfigTester(self       ,     config_class=__snake_case       ,     has_text_modality=__snake_case       ,     hidden_size=37       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											self.config_tester.run_common_tests()


					@unittest.skip(reason='ViT does not use inputs_embeds'       )
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											pass


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





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

											for model_class in self.all_model_classes:
																	_snake_case:   Optional[int]  					=		model_class(__snake_case       )
																	self.assertIsInstance(model.get_input_embeddings()       ,     (nn.Module)       )
																	_snake_case:   Optional[Any]  					=		model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(__snake_case       ,     nn.Linear       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





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

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

																	_snake_case:   str  					=		['pixel_values']
																	self.assertListEqual(arg_names[:1]       ,     __snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   List[Any]  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   int  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_image_modeling(*__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											_snake_case:   str  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__snake_case       )




					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_snake_case:   Any  					=		ViTModel.from_pretrained(__snake_case       )
																	self.assertIsNotNone(__snake_case       )





def       lowercase_      (		)   ->List[Any]:
						_snake_case:   Dict  					=		Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
						return image







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


					@cached_property
					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											return ViTImageProcessor.from_pretrained('google/vit-base-patch16-224'       ) if is_vision_available() else None


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'       ).to(__snake_case       )

											_snake_case:   Dict  					=		self.default_image_processor
											_snake_case:   Optional[Any]  					=		prepare_img()
											_snake_case:   List[str]  					=		image_processor(images=__snake_case       ,     return_tensors='pt'       ).to(__snake_case       )

											# forward pass
											with torch.no_grad():
																	_snake_case:   Optional[int]  					=		model(**__snake_case       )

											# verify the logits
											_snake_case:   Union[str, Any]  					=		torch.Size((1, 10_00)       )
											self.assertEqual(outputs.logits.shape       ,     __snake_case       )

											_snake_case:   Dict  					=		torch.tensor([-0.2_744, 0.8_215, -0.0_836]       ).to(__snake_case       )

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


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											_snake_case:   str  					=		ViTModel.from_pretrained('facebook/dino-vits8'       ).to(__snake_case       )

											_snake_case:   Any  					=		ViTImageProcessor.from_pretrained('facebook/dino-vits8'       ,     size=4_80       )
											_snake_case:   Optional[int]  					=		prepare_img()
											_snake_case:   Dict  					=		image_processor(images=__snake_case       ,     return_tensors='pt'       )
											_snake_case:   Optional[Any]  					=		inputs.pixel_values.to(__snake_case       )

											# forward pass
											with torch.no_grad():
																	_snake_case:   str  					=		model(__snake_case       ,     interpolate_pos_encoding=__snake_case       )

											# verify the logits
											_snake_case:   List[str]  					=		torch.Size((1, 36_01, 3_84)       )
											self.assertEqual(outputs.last_hidden_state.shape       ,     __snake_case       )

											_snake_case:   Any  					=		torch.tensor(
											    [[4.2_340, 4.3_906, -6.6_692], [4.5_463, 1.8_928, -6.7_257], [4.4_429, 0.8_496, -5.8_585]]       ).to(__snake_case       )

											self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]       ,     __snake_case       ,     atol=1e-4       )       )




					@slow
					@require_accelerate
					@require_torch_gpu
					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   List[Any]  					=		ViTModel.from_pretrained('facebook/dino-vits8'       ,     torch_dtype=torch.floataa       ,     device_map='auto'       )
											_snake_case:   Dict  					=		self.default_image_processor

											_snake_case:   Any  					=		prepare_img()
											_snake_case:   str  					=		image_processor(images=__snake_case       ,     return_tensors='pt'       )
											_snake_case:   Any  					=		inputs.pixel_values.to(__snake_case       )

											# forward pass to make sure inference works in fp16
											with torch.no_grad():
																	_snake_case:   int  					=		model(__snake_case       )





