


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 	_a							(      UpperCamelCase__      ):






							"""simple docstring"""




							def 				lowerCamelCase_			(							self:     Dict							)      ->       str:





									"""simple docstring"""







									lowercase__						=				tempfile.mkdtemp()
									lowercase__						=				8

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

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

							def 				lowerCamelCase_			(							self:     Tuple							)      ->       DPRQuestionEncoderTokenizer:





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     Tuple							)      ->       DPRContextEncoderTokenizer:





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     str							)      ->       BartTokenizer:





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     int							)      ->       Tuple:





									"""simple docstring"""







									shutil.rmtree(self.tmpdirname							)

							def 				lowerCamelCase_			(							self:     str							)      ->       Optional[Any]:





									"""simple docstring"""







									lowercase__						=				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 				lowerCamelCase_			(							self:     List[Any]							)      ->       Any:





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     Any   ,						UpperCamelCase_:     bool							)      ->       Union[str, Any]:





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     str							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				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							)

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

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

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

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





									"""simple docstring"""







									lowercase__						=				1
									lowercase__						=				self.get_dummy_canonical_hf_index_retriever()
									lowercase__						=				np.array(
									    [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]   ,						dtype=np.floataa							)
									lowercase__						=				retriever.retrieve(lowerCAmelCase__   ,						n_docs=lowerCAmelCase__							)
									self.assertEqual(retrieved_doc_embeds.shape   ,						(2, n_docs, self.retrieval_vector_size)							)
									self.assertEqual(len(lowerCAmelCase__							)   ,						2							)
									self.assertEqual(sorted(doc_dicts[0]							)   ,						['''embeddings''', '''id''', '''text''', '''title''']							)
									self.assertEqual(len(doc_dicts[0]['''id''']							)   ,						lowerCAmelCase__							)
									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 				lowerCamelCase_			(							self:     Optional[int]							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				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:
													lowercase__						=				self.get_dummy_dataset()
													retriever.save_pretrained(lowerCAmelCase__							)
													lowercase__						=				RagRetriever.from_pretrained(lowerCAmelCase__							)
													self.assertIsInstance(lowerCAmelCase__   ,						lowerCAmelCase__							)
													lowercase__						=				np.array(
													    [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]   ,						dtype=np.floataa							)
													lowercase__						=				retriever.retrieve(lowerCAmelCase__   ,						n_docs=1							)
													self.assertTrue(out is not None							)

							def 				lowerCamelCase_			(							self:     List[str]							)      ->       Optional[Any]:





									"""simple docstring"""







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





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     Dict							)      ->       List[str]:





									"""simple docstring"""







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





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     Any							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				1
									lowercase__						=				self.get_dummy_legacy_index_retriever()
									lowercase__						=				np.array(
									    [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]   ,						dtype=np.floataa							)
									lowercase__						=				retriever.retrieve(lowerCAmelCase__   ,						n_docs=lowerCAmelCase__							)
									self.assertEqual(retrieved_doc_embeds.shape   ,						(2, n_docs, self.retrieval_vector_size)							)
									self.assertEqual(len(lowerCAmelCase__							)   ,						2							)
									self.assertEqual(sorted(doc_dicts[0]							)   ,						['''text''', '''title''']							)
									self.assertEqual(len(doc_dicts[0]['''text''']							)   ,						lowerCAmelCase__							)
									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 				lowerCamelCase_			(							self:     List[str]							)      ->       List[str]:





									"""simple docstring"""







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

							@require_torch
							@require_tokenizers
							@require_sentencepiece
							def 				lowerCamelCase_			(							self:     Any							)      ->       str:





									"""simple docstring"""







									import torch

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

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


							@require_torch
							@require_tokenizers
							@require_sentencepiece
							def 				lowerCamelCase_			(							self:     Optional[int]							)      ->       Optional[Any]:





									"""simple docstring"""







									lowercase__						=				self.get_dpr_ctx_encoder_tokenizer()
									lowercase__						=				1
									lowercase__						=				self.get_dummy_custom_hf_index_retriever(from_disk=lowerCAmelCase__							)
									retriever.set_ctx_encoder_tokenizer(lowerCAmelCase__							)

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

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






def       _a			(				SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		lowercase__						=				len(SCREAMING_SNAKE_CASE					)
		lowercase__						=				[]
		for i in range(len(SCREAMING_SNAKE_CASE					) - pat_len + 1					):
				lowercase__						=				True
				for j in range(SCREAMING_SNAKE_CASE					):
						if s[i + j] != pattern[j]:
								lowercase__						=				False
								break
				if match_found:
						position.append(SCREAMING_SNAKE_CASE					)
		return position


if __name__ == "__main__":
			assert naive_pattern_search('ABCDEFG', 'DE') == [3]
			print(naive_pattern_search('ABAAABCDBBABCDDEBCABC', 'ABC'))


