



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.









def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			int(SCREAMING_SNAKE_CASE_					)
			if decimal in (0, 1):  # Exit cases for the recursion
						return str(SCREAMING_SNAKE_CASE_					)
			lowercase__					,    lowercase__   					=			divmod(SCREAMING_SNAKE_CASE_  ,					2					)
			return binary_recursive(SCREAMING_SNAKE_CASE_					) + str(SCREAMING_SNAKE_CASE_					)



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			str(SCREAMING_SNAKE_CASE_					).strip()
			if not number:
						raise ValueError("No input value was provided"					)
			lowercase__   					=			"-" if number.startswith("-"					) else ""
			lowercase__   					=			number.lstrip("-"					)
			if not number.isnumeric():
						raise ValueError("Input value is not an integer"					)
			return f'''{negative}0b{binary_recursive(int(SCREAMING_SNAKE_CASE_					)					)}'''


if __name__ == "__main__":
							from doctest import testmod

							testmod()




