




import argparse

import requests
import torch
from PIL import Image

from transformers import SwinConfig, SwinForMaskedImageModeling, ViTImageProcessor



def       UpperCAmelCase       (     _snake_case							):
							lowerCAmelCase   =		SwinConfig(image_size=192							)

							if "base" in model_name:
														lowerCAmelCase   =		6
														lowerCAmelCase   =		128
														lowerCAmelCase   =		(2, 2, 18, 2)
														lowerCAmelCase   =		(4, 8, 16, 32)
							elif "large" in model_name:
														lowerCAmelCase   =		12
														lowerCAmelCase   =		192
														lowerCAmelCase   =		(2, 2, 18, 2)
														lowerCAmelCase   =		(6, 12, 24, 48)
							else:
														raise ValueError('''Model not supported, only supports base and large variants'''							)

							lowerCAmelCase   =		window_size
							lowerCAmelCase   =		embed_dim
							lowerCAmelCase   =		depths
							lowerCAmelCase   =		num_heads

							return config



def       UpperCAmelCase       (     _snake_case							):



							if "encoder.mask_token" in name:
														lowerCAmelCase   =		name.replace('''encoder.mask_token'''						, '''embeddings.mask_token'''							)
							if "encoder.patch_embed.proj" in name:
														lowerCAmelCase   =		name.replace('''encoder.patch_embed.proj'''						, '''embeddings.patch_embeddings.projection'''							)
							if "encoder.patch_embed.norm" in name:
														lowerCAmelCase   =		name.replace('''encoder.patch_embed.norm'''						, '''embeddings.norm'''							)
							if "attn.proj" in name:
														lowerCAmelCase   =		name.replace('''attn.proj'''						, '''attention.output.dense'''							)
							if "attn" in name:
														lowerCAmelCase   =		name.replace('''attn'''						, '''attention.self'''							)
							if "norm1" in name:
														lowerCAmelCase   =		name.replace('''norm1'''						, '''layernorm_before'''							)
							if "norm2" in name:
														lowerCAmelCase   =		name.replace('''norm2'''						, '''layernorm_after'''							)
							if "mlp.fc1" in name:
														lowerCAmelCase   =		name.replace('''mlp.fc1'''						, '''intermediate.dense'''							)
							if "mlp.fc2" in name:
														lowerCAmelCase   =		name.replace('''mlp.fc2'''						, '''output.dense'''							)

							if name == "encoder.norm.weight":
														lowerCAmelCase   =		'''layernorm.weight'''
							if name == "encoder.norm.bias":
														lowerCAmelCase   =		'''layernorm.bias'''

							if "decoder" in name:
														pass
							else:
														lowerCAmelCase   =		'''swin.''' + name

							return name



def       UpperCAmelCase       (     _snake_case						, _snake_case							):



							for key in orig_state_dict.copy().keys():
														lowerCAmelCase   =		orig_state_dict.pop(_SCREAMING_SNAKE_CASE							)

														if "attn_mask" in key:
																					pass
														elif "qkv" in key:
																					lowerCAmelCase   =		key.split('''.'''							)
																					lowerCAmelCase   =		int(key_split[2]							)
																					lowerCAmelCase   =		int(key_split[4]							)
																					lowerCAmelCase   =		model.swin.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

																					if "weight" in key:
																												lowerCAmelCase   =		val[:dim, :]
																												lowerCAmelCase   =		val[
																												    dim : dim * 2, :
																												]
																												lowerCAmelCase   =		val[-dim:, :]
																					else:
																												lowerCAmelCase   =		val[
																												    :dim
																												]
																												lowerCAmelCase   =		val[
																												    dim : dim * 2
																												]
																												lowerCAmelCase   =		val[
																												    -dim:
																												]
														else:
																					lowerCAmelCase   =		val

							return orig_state_dict



def       UpperCAmelCase       (     _snake_case						, _snake_case						, _snake_case						, _snake_case							):
							lowerCAmelCase   =		torch.load(_SCREAMING_SNAKE_CASE						, map_location='''cpu'''							)['''model''']

							lowerCAmelCase   =		get_swin_config(_SCREAMING_SNAKE_CASE							)
							lowerCAmelCase   =		SwinForMaskedImageModeling(_SCREAMING_SNAKE_CASE							)
							model.eval()

							lowerCAmelCase   =		convert_state_dict(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)
							model.load_state_dict(_SCREAMING_SNAKE_CASE							)

							lowerCAmelCase   =		'''http://images.cocodataset.org/val2017/000000039769.jpg'''

							lowerCAmelCase   =		ViTImageProcessor(size={'''height''': 192, '''width''': 192}							)
							lowerCAmelCase   =		Image.open(requests.get(_SCREAMING_SNAKE_CASE						, stream=_SCREAMING_SNAKE_CASE							).raw							)
							lowerCAmelCase   =		image_processor(images=_SCREAMING_SNAKE_CASE						, return_tensors='''pt'''							)

							with torch.no_grad():
														lowerCAmelCase   =		model(**_SCREAMING_SNAKE_CASE							).logits

							print(outputs.keys()							)
							print('''Looks ok!'''							)

							if pytorch_dump_folder_path is not None:
														print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""							)
														model.save_pretrained(_SCREAMING_SNAKE_CASE							)

														print(F"""Saving image processor to {pytorch_dump_folder_path}"""							)
														image_processor.save_pretrained(_SCREAMING_SNAKE_CASE							)

							if push_to_hub:
														print(F"""Pushing model and image processor for {model_name} to hub"""							)
														model.push_to_hub(F"""microsoft/{model_name}"""							)
														image_processor.push_to_hub(F"""microsoft/{model_name}"""							)


if __name__ == "__main__":
		UpperCAmelCase_						=argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--model_name""",
		    default="""swin-base-simmim-window6-192""",
		    type=str,
		    choices=["""swin-base-simmim-window6-192""", """swin-large-simmim-window12-192"""],
		    help="""Name of the Swin SimMIM model you'd like to convert.""",
		)
		parser.add_argument(
		    """--checkpoint_path""",
		    default="""/Users/nielsrogge/Documents/SwinSimMIM/simmim_pretrain__swin_base__img192_window6__100ep.pth""",
		    type=str,
		    help="""Path to the original PyTorch checkpoint (.pth file).""",
		)
		parser.add_argument(
		    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
		)
		parser.add_argument(
		    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
		)

		UpperCAmelCase_						=parser.parse_args()
		convert_swin_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)









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 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				def       __snake_case		( self						):
											lowerCAmelCase   =		tempfile.mkdtemp()
											lowerCAmelCase   =		8

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

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



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



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



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



				def       __snake_case		( self						):
											shutil.rmtree(self.tmpdirname						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		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       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dummy_dataset()
											lowerCAmelCase   =		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:
																		lowerCAmelCase   =		dataset
																		lowerCAmelCase   =		RagRetriever(
																		    UpperCAmelCase_       , question_encoder_tokenizer=self.get_dpr_tokenizer()       , generator_tokenizer=self.get_bart_tokenizer()       , )
											return retriever



				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.get_dummy_dataset()
											lowerCAmelCase   =		RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       , question_encoder=DPRConfig().to_dict()       , generator=BartConfig().to_dict()       , index_name='''custom'''       , )
											if from_disk:
																		lowerCAmelCase   =		os.path.join(self.tmpdirname       , '''dataset'''						)
																		lowerCAmelCase   =		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
																		lowerCAmelCase   =		RagRetriever(
																		    UpperCAmelCase_       , question_encoder_tokenizer=self.get_dpr_tokenizer()       , generator_tokenizer=self.get_bart_tokenizer()       , )
											else:
																		lowerCAmelCase   =		RagRetriever(
																		    UpperCAmelCase_       , question_encoder_tokenizer=self.get_dpr_tokenizer()       , generator_tokenizer=self.get_bart_tokenizer()       , index=CustomHFIndex(config.retrieval_vector_size       , UpperCAmelCase_						)       , )
											return retriever



				def       __snake_case		( self						):
											lowerCAmelCase   =		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						)

											lowerCAmelCase   =		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'''						)						)

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

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



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_canonical_hf_index_retriever()
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''embeddings''', '''id''', '''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''id''']						)       , UpperCAmelCase_						)
											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       __snake_case		( self						):
											lowerCAmelCase   =		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:
																									lowerCAmelCase   =		self.get_dummy_dataset()
																									retriever.save_pretrained(UpperCAmelCase_						)
																									lowerCAmelCase   =		RagRetriever.from_pretrained(UpperCAmelCase_						)
																									self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_						)
																									lowerCAmelCase   =		np.array(
																									    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
																									lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=1						)
																									self.assertTrue(out is not None						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''embeddings''', '''id''', '''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''id''']						)       , UpperCAmelCase_						)
											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       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(UpperCAmelCase_						)
																		lowerCAmelCase   =		RagRetriever.from_pretrained(UpperCAmelCase_						)
																		self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_						)
																		lowerCAmelCase   =		np.array(
																		    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
																		lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=1						)
																		self.assertTrue(out is not None						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''embeddings''', '''id''', '''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''id''']						)       , UpperCAmelCase_						)
											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       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(UpperCAmelCase_						)
																		lowerCAmelCase   =		RagRetriever.from_pretrained(UpperCAmelCase_						)
																		self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_						)
																		lowerCAmelCase   =		np.array(
																		    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
																		lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=1						)
																		self.assertTrue(out is not None						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_legacy_index_retriever()
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''text''']						)       , UpperCAmelCase_						)
											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       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dummy_legacy_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(UpperCAmelCase_						)
																		lowerCAmelCase   =		RagRetriever.from_pretrained(UpperCAmelCase_						)
																		self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_						)
																		lowerCAmelCase   =		np.array(
																		    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
																		lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=1						)
																		self.assertTrue(out is not None						)



				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def       __snake_case		( self						):
											import torch

											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_canonical_hf_index_retriever()
											lowerCAmelCase   =		[[5, 7], [10, 11]]
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   =		retriever(UpperCAmelCase_       , UpperCAmelCase_       , prefix=retriever.config.generator.prefix       , n_docs=UpperCAmelCase_						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		(
											    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(UpperCAmelCase_       , UpperCAmelCase_						)
											self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_						)
											self.assertIsInstance(UpperCAmelCase_       , np.ndarray						)

											lowerCAmelCase   =		retriever(
											    UpperCAmelCase_       , UpperCAmelCase_       , prefix=retriever.config.generator.prefix       , n_docs=UpperCAmelCase_       , return_tensors='''pt'''       , )
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		(  # 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(UpperCAmelCase_       , torch.Tensor						)
											self.assertIsInstance(UpperCAmelCase_       , torch.Tensor						)
											self.assertIsInstance(UpperCAmelCase_       , torch.Tensor						)






				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dpr_ctx_encoder_tokenizer()
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											retriever.set_ctx_encoder_tokenizer(UpperCAmelCase_						)

											lowerCAmelCase   =		[[5, 7], [10, 11]]
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   =		retriever(UpperCAmelCase_       , UpperCAmelCase_       , prefix=retriever.config.generator.prefix       , n_docs=UpperCAmelCase_						)

											self.assertEqual(
											    len(UpperCAmelCase_						)       , 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''')						)       , UpperCAmelCase_						)  # check for doc token related keys in dictionary.



