



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 			_snake_case     (   lowercase__):




				def 							A__ (  self   :  Tuple  ):
							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(__lowercase,   exist_ok=__lowercase  )
							lowercase__   					=			os.path.join(__lowercase,   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(__lowercase,   range(len(__lowercase  )  )  )  )
							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(__lowercase,   exist_ok=__lowercase  )
							lowercase__   					=			os.path.join(__lowercase,   BART_VOCAB_FILES_NAMES["vocab_file"]  )
							lowercase__   					=			os.path.join(__lowercase,   BART_VOCAB_FILES_NAMES["merges_file"]  )
							with open(self.vocab_file,   "w",   encoding="utf-8"  ) as fp:
										fp.write(json.dumps(__lowercase  ) + "\n"  )
							with open(self.merges_file,   "w",   encoding="utf-8"  ) as fp:
										fp.write("\n".join(__lowercase  )  )




				def 							A__ (  self   :  Tuple  ):
							return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname,   "dpr_tokenizer"  )  )




				def 							A__ (  self   :  Tuple  ):
							return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname,   "dpr_tokenizer"  )  )




				def 							A__ (  self   :  str  ):
							return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname,   "bart_tokenizer"  )  )




				def 							A__ (  self   :  Any  ):
							shutil.rmtree(self.tmpdirname  )




				def 							A__ (  self   :  List[str]  ):
							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 							A__ (  self   :  str  ):
							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(
										    __lowercase,   question_encoder_tokenizer=self.get_dpr_tokenizer(),   generator_tokenizer=self.get_bart_tokenizer(),   )
							return retriever




				def 							A__ (  self   :  Any,   __lowercase   :  bool  ):
							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(
										    __lowercase,   question_encoder_tokenizer=self.get_dpr_tokenizer(),   generator_tokenizer=self.get_bart_tokenizer(),   )
							else:
										lowercase__   					=			RagRetriever(
										    __lowercase,   question_encoder_tokenizer=self.get_dpr_tokenizer(),   generator_tokenizer=self.get_bart_tokenizer(),   index=CustomHFIndex(config.retrieval_vector_size,   __lowercase  ),   )
							return retriever




				def 							A__ (  self   :  Optional[int]  ):
							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(__lowercase,   open(__lowercase,   "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(
							    __lowercase,   question_encoder_tokenizer=self.get_dpr_tokenizer(),   generator_tokenizer=self.get_bart_tokenizer()  )
							return retriever




				def 							A__ (  self   :  Union[str, Any]  ):
							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__					,    lowercase__					,    lowercase__   					=			retriever.retrieve(__lowercase,   n_docs=__lowercase  )
							self.assertEqual(retrieved_doc_embeds.shape,   (2, n_docs, self.retrieval_vector_size)  )
							self.assertEqual(len(__lowercase  ),   2  )
							self.assertEqual(sorted(doc_dicts[0]  ),   ["embeddings", "id", "text", "title"]  )
							self.assertEqual(len(doc_dicts[0]["id"]  ),   __lowercase  )
							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 							A__ (  self   :  int  ):
							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(__lowercase  )
													lowercase__   					=			RagRetriever.from_pretrained(__lowercase  )
													self.assertIsInstance(__lowercase,   __lowercase  )
													lowercase__   					=			np.array(
													    [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )],   dtype=np.floataa  )
													lowercase__   					=			retriever.retrieve(__lowercase,   n_docs=1  )
													self.assertTrue(out is not None  )




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




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




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




				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def 							A__ (  self   :  Dict  ):
							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(__lowercase,   __lowercase,   prefix=retriever.config.generator.prefix,   n_docs=__lowercase  )
							lowercase__					,    lowercase__					,    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(__lowercase,   __lowercase  )
							self.assertIsInstance(__lowercase,   __lowercase  )
							self.assertIsInstance(__lowercase,   np.ndarray  )

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




				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def 							A__ (  self   :  Any  ):
							lowercase__   					=			self.get_dpr_ctx_encoder_tokenizer()
							lowercase__   					=			1
							lowercase__   					=			self.get_dummy_custom_hf_index_retriever(from_disk=__lowercase  )
							retriever.set_ctx_encoder_tokenizer(__lowercase  )

							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(__lowercase,   __lowercase,   prefix=retriever.config.generator.prefix,   n_docs=__lowercase  )

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









import os
from collections.abc import Iterator



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_ = "."					):
			for dir_path, dir_names, filenames in os.walk(SCREAMING_SNAKE_CASE_					):
						lowercase__   					=			[d for d in dir_names if d != "scripts" and d[0] not in "._"]
						for filename in filenames:
									if filename == "__init__.py":
												continue
									if os.path.splitext(SCREAMING_SNAKE_CASE_					)[1] in (".py", ".ipynb"):
												yield os.path.join(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					).lstrip("./"					)



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_					):
			return f'''{i * "  "}*''' if i else "\n##"



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			old_path.split(os.sep					)
			for i, new_part in enumerate(new_path.split(os.sep					)					):
						if (i + 1 > len(SCREAMING_SNAKE_CASE_					) or old_parts[i] != new_part) and new_part:
									print(f'''{md_prefix(SCREAMING_SNAKE_CASE_					)} {new_part.replace("_"  ,					" "					).title()}'''					)
			return new_path



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_ = "."					):
			lowercase__   					=			""
			for filepath in sorted(good_file_paths(SCREAMING_SNAKE_CASE_					)					):
						lowercase__					,    lowercase__   					=			os.path.split(SCREAMING_SNAKE_CASE_					)
						if filepath != old_path:
									lowercase__   					=			print_path(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
						lowercase__   					=			(filepath.count(os.sep					) + 1) if filepath else 0
						lowercase__   					=			f'''{filepath}/{filename}'''.replace(" "  ,					"%20"					)
						lowercase__   					=			os.path.splitext(filename.replace("_"  ,					" "					).title()					)[0]
						print(f'''{md_prefix(SCREAMING_SNAKE_CASE_					)} [{filename}]({url})'''					)


if __name__ == "__main__":
							print_directory_md(""".""")




