






"""simple docstring"""

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 		__lowerCamelCase				( SCREAMING_SNAKE_CASE__			):





		'''simple docstring'''






		def    lowerCamelCase       (   self							:	str	):
				lowerCAmelCase_   :							List[Any]									=			tempfile.mkdtemp()
				lowerCAmelCase_   :							List[Any]									=			8

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

				lowerCAmelCase_   :							Optional[Any]									=			os.path.join(self.tmpdirname						,			"bart_tokenizer"	)
				os.makedirs(a_						,			exist_ok=a_	)
				lowerCAmelCase_   :							Tuple									=			os.path.join(a_						,			BART_VOCAB_FILES_NAMES["vocab_file"]	)
				lowerCAmelCase_   :							Union[str, Any]									=			os.path.join(a_						,			BART_VOCAB_FILES_NAMES["merges_file"]	)
				with open(self.vocab_file						,			"w"						,			encoding="utf-8"	) as fp:
						fp.write(json.dumps(a_	) + "\n"	)
				with open(self.merges_file						,			"w"						,			encoding="utf-8"	) as fp:
						fp.write("\n".join(a_	)	)





		def    lowerCamelCase       (   self							:	List[Any]	):
				return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname						,			"dpr_tokenizer"	)	)





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





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





		def    lowerCamelCase       (   self							:	List[Any]	):
				shutil.rmtree(self.tmpdirname	)





		def    lowerCamelCase       (   self							:	List[Any]	):
				lowerCAmelCase_   :							Optional[int]									=			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[str]	):
				lowerCAmelCase_   :							Any									=			self.get_dummy_dataset()
				lowerCAmelCase_   :							Union[str, Any]									=			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_   :							Any									=			dataset
						lowerCAmelCase_   :							str									=			RagRetriever(
						    a_						,			question_encoder_tokenizer=self.get_dpr_tokenizer()						,			generator_tokenizer=self.get_bart_tokenizer()						,			)
				return retriever





		def    lowerCamelCase       (   self							:	Union[str, Any]						,			a_							:	bool	):
				lowerCAmelCase_   :							int									=			self.get_dummy_dataset()
				lowerCAmelCase_   :							Optional[int]									=			RagConfig(
				    retrieval_vector_size=self.retrieval_vector_size						,			question_encoder=DPRConfig().to_dict()						,			generator=BartConfig().to_dict()						,			index_name="custom"						,			)
				if from_disk:
						lowerCAmelCase_   :							Dict									=			os.path.join(self.tmpdirname						,			"dataset"	)
						lowerCAmelCase_   :							Optional[int]									=			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_   :							List[Any]									=			RagRetriever(
						    a_						,			question_encoder_tokenizer=self.get_dpr_tokenizer()						,			generator_tokenizer=self.get_bart_tokenizer()						,			)
				else:
						lowerCAmelCase_   :							str									=			RagRetriever(
						    a_						,			question_encoder_tokenizer=self.get_dpr_tokenizer()						,			generator_tokenizer=self.get_bart_tokenizer()						,			index=CustomHFIndex(config.retrieval_vector_size						,			a_	)						,			)
				return retriever





		def    lowerCamelCase       (   self							:	List[Any]	):
				lowerCAmelCase_   :							Optional[Any]									=			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_   :							Union[str, Any]									=			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_   :							List[Any]									=			os.path.join(self.tmpdirname						,			"psgs_w100.tsv.pkl"	)
				lowerCAmelCase_   :							List[Any]									=			{sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
				pickle.dump(a_						,			open(a_						,			"wb"	)	)

				lowerCAmelCase_   :							Dict									=			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_   :							List[str]									=			RagRetriever(
				    a_						,			question_encoder_tokenizer=self.get_dpr_tokenizer()						,			generator_tokenizer=self.get_bart_tokenizer()	)
				return retriever





		def    lowerCamelCase       (   self							:	Tuple	):
				lowerCAmelCase_   :							Optional[Any]									=			1
				lowerCAmelCase_   :							List[Any]									=			self.get_dummy_canonical_hf_index_retriever()
				lowerCAmelCase_   :							Tuple									=			np.array(
				    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]						,			dtype=np.floataa	)
				lowerCAmelCase_				,	lowerCAmelCase_				,	lowerCAmelCase_   :							Union[str, Any]									=			retriever.retrieve(a_						,			n_docs=a_	)
				self.assertEqual(retrieved_doc_embeds.shape						,			(2, n_docs, self.retrieval_vector_size)	)
				self.assertEqual(len(a_	)						,			2	)
				self.assertEqual(sorted(doc_dicts[0]	)						,			["embeddings", "id", "text", "title"]	)
				self.assertEqual(len(doc_dicts[0]["id"]	)						,			a_	)
				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							:	Dict	):
				lowerCAmelCase_   :							Optional[int]									=			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_   :							Tuple									=			self.get_dummy_dataset()
								retriever.save_pretrained(a_	)
								lowerCAmelCase_   :							Union[str, Any]									=			RagRetriever.from_pretrained(a_	)
								self.assertIsInstance(a_						,			a_	)
								lowerCAmelCase_   :							str									=			np.array(
								    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]						,			dtype=np.floataa	)
								lowerCAmelCase_   :							Dict									=			retriever.retrieve(a_						,			n_docs=1	)
								self.assertTrue(out is not None	)





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





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





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





		@require_torch
		@require_tokenizers
		@require_sentencepiece
		def    lowerCamelCase       (   self							:	Optional[Any]	):
				import torch

				lowerCAmelCase_   :							Union[str, Any]									=			1
				lowerCAmelCase_   :							Tuple									=			self.get_dummy_canonical_hf_index_retriever()
				lowerCAmelCase_   :							List[Any]									=			[[5, 7], [10, 11]]
				lowerCAmelCase_   :							Optional[Any]									=			np.array(
				    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]						,			dtype=np.floataa	)
				lowerCAmelCase_   :							Tuple									=			retriever(a_						,			a_						,			prefix=retriever.config.generator.prefix						,			n_docs=a_	)
				lowerCAmelCase_				,	lowerCAmelCase_				,	lowerCAmelCase_   :							Any									=			(
				    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(a_						,			a_	)
				self.assertIsInstance(a_						,			a_	)
				self.assertIsInstance(a_						,			np.ndarray	)

				lowerCAmelCase_   :							int									=			retriever(
				    a_						,			a_						,			prefix=retriever.config.generator.prefix						,			n_docs=a_						,			return_tensors="pt"						,			)
				lowerCAmelCase_				,	lowerCAmelCase_				,	lowerCAmelCase_				,	lowerCAmelCase_   :							Optional[Any]									=			(  # 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(a_						,			torch.Tensor	)
				self.assertIsInstance(a_						,			torch.Tensor	)
				self.assertIsInstance(a_						,			torch.Tensor	)





		@require_torch
		@require_tokenizers
		@require_sentencepiece
		def    lowerCamelCase       (   self							:	Optional[Any]	):
				lowerCAmelCase_   :							str									=			self.get_dpr_ctx_encoder_tokenizer()
				lowerCAmelCase_   :							Dict									=			1
				lowerCAmelCase_   :							str									=			self.get_dummy_custom_hf_index_retriever(from_disk=a_	)
				retriever.set_ctx_encoder_tokenizer(a_	)

				lowerCAmelCase_   :							List[str]									=			[[5, 7], [10, 11]]
				lowerCAmelCase_   :							List[str]									=			np.array(
				    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]						,			dtype=np.floataa	)
				lowerCAmelCase_   :							List[str]									=			retriever(a_						,			a_						,			prefix=retriever.config.generator.prefix						,			n_docs=a_	)

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

def    a__       (						__UpperCamelCase ):
      if not head:
            return True
      # split the list to two parts
      SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     head.next, head
      while fast and fast.next:
            SCREAMING_SNAKE_CASE_      =     fast.next.next
            SCREAMING_SNAKE_CASE_      =     slow.next
      SCREAMING_SNAKE_CASE_      =     slow.next
      SCREAMING_SNAKE_CASE_      =     None  # Don't forget here! But forget still works!
      # reverse the second part
      SCREAMING_SNAKE_CASE_      =     None
      while second:
            SCREAMING_SNAKE_CASE_      =     second.next
            SCREAMING_SNAKE_CASE_      =     node
            SCREAMING_SNAKE_CASE_      =     second
            SCREAMING_SNAKE_CASE_      =     nxt
      # compare two parts
      # second part has the same or one less node
      while node:
            if node.val != head.val:
                  return False
            SCREAMING_SNAKE_CASE_      =     node.next
            SCREAMING_SNAKE_CASE_      =     head.next
      return True


def    a__       (						__UpperCamelCase ):
      if not head or not head.next:
            return True

      # 1. Get the midpoint (slow)
      SCREAMING_SNAKE_CASE_      =     SCREAMING_SNAKE_CASE_      =     SCREAMING_SNAKE_CASE_      =     head
      while fast and fast.next:
            SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     fast.next.next, slow.next

      # 2. Push the second half into the stack
      SCREAMING_SNAKE_CASE_      =     [slow.val]
      while slow.next:
            SCREAMING_SNAKE_CASE_      =     slow.next
            stack.append(slow.val )

      # 3. Comparison
      while stack:
            if stack.pop() != cur.val:
                  return False
            SCREAMING_SNAKE_CASE_      =     cur.next

      return True





def    a__       (						__UpperCamelCase ):
      if not head or not head.next:
            return True
      SCREAMING_SNAKE_CASE_      =     {}
      SCREAMING_SNAKE_CASE_      =     0
      while head:
            if head.val in d:
                  d[head.val].append(__UpperCamelCase )
            else:
                  SCREAMING_SNAKE_CASE_      =     [pos]
            SCREAMING_SNAKE_CASE_      =     head.next
            pos += 1
      SCREAMING_SNAKE_CASE_      =     pos - 1
      SCREAMING_SNAKE_CASE_      =     0
      for v in d.values():
            if len(__UpperCamelCase ) % 2 != 0:
                  middle += 1
            else:
                  SCREAMING_SNAKE_CASE_      =     0
                  for i in range(0   ,					len(__UpperCamelCase ) ):
                        if v[i] + v[len(__UpperCamelCase ) - 1 - step] != checksum:
                              return False
                        step += 1
            if middle > 1:
                  return False
      return True


