




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.









# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from typing import Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import randn_tensor
from .scheduling_utils import SchedulerMixin


class 							__UpperCamelCase			( __UpperCAmelCase     ,					__UpperCAmelCase  ):





				'''simple docstring'''



				__a							: Any										=1



				@register_to_config
				def __init__( self       , UpperCAmelCase_=20_00       , UpperCAmelCase_=0.1       , UpperCAmelCase_=20       , UpperCAmelCase_=1E-3						):
											lowerCAmelCase   =		None
											lowerCAmelCase   =		None
											lowerCAmelCase   =		None



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_ = None						):
											lowerCAmelCase   =		torch.linspace(1       , self.config.sampling_eps       , UpperCAmelCase_       , device=UpperCAmelCase_						)



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_=None						):
											if self.timesteps is None:
																		raise ValueError(
																		    '''`self.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler'''						)

											# TODO(Patrick) better comments + non-PyTorch
											# postprocess model score
											lowerCAmelCase   =		(
											    -0.25 * t**2 * (self.config.beta_max - self.config.beta_min) - 0.5 * t * self.config.beta_min
											)
											lowerCAmelCase   =		torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff						)						)
											lowerCAmelCase   =		std.flatten()
											while len(std.shape						) < len(score.shape						):
																		lowerCAmelCase   =		std.unsqueeze(-1						)
											lowerCAmelCase   =		-score / std

											# compute
											lowerCAmelCase   =		-1.0 / len(self.timesteps						)

											lowerCAmelCase   =		self.config.beta_min + t * (self.config.beta_max - self.config.beta_min)
											lowerCAmelCase   =		beta_t.flatten()
											while len(beta_t.shape						) < len(x.shape						):
																		lowerCAmelCase   =		beta_t.unsqueeze(-1						)
											lowerCAmelCase   =		-0.5 * beta_t * x

											lowerCAmelCase   =		torch.sqrt(UpperCAmelCase_						)
											lowerCAmelCase   =		drift - diffusion**2 * score
											lowerCAmelCase   =		x + drift * dt

											# add noise
											lowerCAmelCase   =		randn_tensor(x.shape       , layout=x.layout       , generator=UpperCAmelCase_       , device=x.device       , dtype=x.dtype						)
											lowerCAmelCase   =		x_mean + diffusion * math.sqrt(-dt						) * noise

											return x, x_mean






				def __len__( self						):
											return self.config.num_train_timesteps



