





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   _SCREAMING_SNAKE_CASE		(      __SCREAMING_SNAKE_CASE       ):







 '''simple docstring'''

 def 				SCREAMING_SNAKE_CASE_   (self							:  int)			->Dict:



   '''simple docstring'''



   lowerCamelCase__: Any         =tempfile.mkdtemp()
   lowerCamelCase__: Any         =8

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

   lowerCamelCase__: str         =os.path.join(self.tmpdirname			,					"bart_tokenizer")
   os.makedirs(UpperCAmelCase_			,					exist_ok=UpperCAmelCase_)
   lowerCamelCase__: Dict         =os.path.join(UpperCAmelCase_			,					BART_VOCAB_FILES_NAMES["vocab_file"])
   lowerCamelCase__: List[str]         =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 				SCREAMING_SNAKE_CASE_   (self							:  Optional[Any])			->DPRQuestionEncoderTokenizer:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any])			->DPRContextEncoderTokenizer:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Optional[int])			->BartTokenizer:



   '''simple docstring'''



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

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



   '''simple docstring'''



   shutil.rmtree(self.tmpdirname)

 def 				SCREAMING_SNAKE_CASE_   (self							:  int)			->Tuple:



   '''simple docstring'''



   lowerCamelCase__: Union[str, Any]         =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 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any])			->int:



   '''simple docstring'''



   lowerCamelCase__: Union[str, Any]         =self.get_dummy_dataset()
   lowerCamelCase__: int         =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__: int         =dataset
     lowerCamelCase__: str         =RagRetriever(
         UpperCAmelCase_			,					question_encoder_tokenizer=self.get_dpr_tokenizer()			,					generator_tokenizer=self.get_bart_tokenizer()			,					)
   return retriever

 def 				SCREAMING_SNAKE_CASE_   (self							:  int			,					UpperCAmelCase_							:  bool)			->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__: int         =self.get_dummy_dataset()
   lowerCamelCase__: 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__: Optional[Any]         =os.path.join(self.tmpdirname			,					"dataset")
     lowerCamelCase__: Tuple         =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__: Dict         =RagRetriever(
         UpperCAmelCase_			,					question_encoder_tokenizer=self.get_dpr_tokenizer()			,					generator_tokenizer=self.get_bart_tokenizer()			,					)
   else:
     lowerCamelCase__: Tuple         =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 				SCREAMING_SNAKE_CASE_   (self							:  Optional[Any])			->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__: List[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__: Dict         =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__: Dict         =os.path.join(self.tmpdirname			,					"psgs_w100.tsv.pkl")
   lowerCamelCase__: str         ={sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
   pickle.dump(UpperCAmelCase_			,					open(UpperCAmelCase_			,					"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__: Any         =RagRetriever(
       UpperCAmelCase_			,					question_encoder_tokenizer=self.get_dpr_tokenizer()			,					generator_tokenizer=self.get_bart_tokenizer())
   return retriever

 def 				SCREAMING_SNAKE_CASE_   (self							:  int)			->List[str]:



   '''simple docstring'''



   lowerCamelCase__: Optional[Any]         =1
   lowerCamelCase__: List[Any]         =self.get_dummy_canonical_hf_index_retriever()
   lowerCamelCase__: Union[str, Any]         =np.array(
       [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]			,					dtype=np.floataa)
   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__: Optional[int]         =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 				SCREAMING_SNAKE_CASE_   (self							:  List[str])			->Any:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Any)			->str:



   '''simple docstring'''



   lowerCamelCase__: Optional[int]         =1
   lowerCamelCase__: Optional[Any]         =self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_)
   lowerCamelCase__: Any         =np.array(
       [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]			,					dtype=np.floataa)
   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__: str         =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 				SCREAMING_SNAKE_CASE_   (self							:  str)			->List[Any]:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Optional[Any])			->Any:



   '''simple docstring'''



   lowerCamelCase__: Tuple         =1
   lowerCamelCase__: Dict         =self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_)
   lowerCamelCase__: str         =np.array(
       [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]			,					dtype=np.floataa)
   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__: Optional[Any]         =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 				SCREAMING_SNAKE_CASE_   (self							:  Optional[int])			->Optional[int]:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  List[Any])			->List[Any]:



   '''simple docstring'''



   lowerCamelCase__: Any         =1
   lowerCamelCase__: Optional[int]         =self.get_dummy_legacy_index_retriever()
   lowerCamelCase__: Any         =np.array(
       [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]			,					dtype=np.floataa)
   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__: Optional[int]         =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 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->List[Any]:



   '''simple docstring'''



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

 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def 				SCREAMING_SNAKE_CASE_   (self							:  Dict)			->Union[str, Any]:



   '''simple docstring'''



   import torch

   lowerCamelCase__: Union[str, Any]         =1
   lowerCamelCase__: Optional[Any]         =self.get_dummy_canonical_hf_index_retriever()
   lowerCamelCase__: 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__: Optional[int]         =retriever(UpperCAmelCase_			,					UpperCAmelCase_			,					prefix=retriever.config.generator.prefix			,					n_docs=UpperCAmelCase_)
   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(UpperCAmelCase_			,					UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_			,					UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_			,					np.ndarray)

   lowerCamelCase__: Optional[Any]         =retriever(
       UpperCAmelCase_			,					UpperCAmelCase_			,					prefix=retriever.config.generator.prefix			,					n_docs=UpperCAmelCase_			,					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(UpperCAmelCase_			,					torch.Tensor)
   self.assertIsInstance(UpperCAmelCase_			,					torch.Tensor)
   self.assertIsInstance(UpperCAmelCase_			,					torch.Tensor)







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



   '''simple docstring'''



   lowerCamelCase__: Any         =self.get_dpr_ctx_encoder_tokenizer()
   lowerCamelCase__: Dict         =1
   lowerCamelCase__: Tuple         =self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_)
   retriever.set_ctx_encoder_tokenizer(UpperCAmelCase_)

   lowerCamelCase__: List[Any]         =[[5, 7], [10, 11]]
   lowerCamelCase__: Any         =np.array(
       [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]			,					dtype=np.floataa)
   lowerCamelCase__: List[Any]         =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.








from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__UpperCamelCase:	Optional[int] = logging.get_logger(__name__)

__UpperCamelCase:	Optional[int] = {
    """google/mobilenet_v1_1.0_224""": """https://huggingface.co/google/mobilenet_v1_1.0_224/resolve/main/config.json""",
    """google/mobilenet_v1_0.75_192""": """https://huggingface.co/google/mobilenet_v1_0.75_192/resolve/main/config.json""",
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
}


class       __lowerCAmelCase					(       _UpperCamelCase      ):


	'''simple docstring'''





	_A													=				"mobilenet_v1"

	def __init__(						self:						Dict,					lowerCamelCase_:						Any=3,					lowerCamelCase_:						Dict=224,					lowerCamelCase_:						int=1.0,					lowerCamelCase_:						List[Any]=8,					lowerCamelCase_:						int="relu6",					lowerCamelCase_:						Optional[int]=True,					lowerCamelCase_:						Any=0.9_9_9,					lowerCamelCase_:						List[str]=0.0_2,					lowerCamelCase_:						Optional[int]=0.0_0_1,					**lowerCamelCase_:						Tuple,					):
			super().__init__(**lowerCamelCase_			)

			if depth_multiplier <= 0:
					raise ValueError('depth_multiplier must be greater than zero.'			)

			lowercase__					:      Optional[int]      =		num_channels
			lowercase__					:      List[Any]      =		image_size
			lowercase__					:      str      =		depth_multiplier
			lowercase__					:      Optional[Any]      =		min_depth
			lowercase__					:      int      =		hidden_act
			lowercase__					:      Dict      =		tf_padding
			lowercase__					:      Optional[Any]      =		classifier_dropout_prob
			lowercase__					:      int      =		initializer_range
			lowercase__					:      Any      =		layer_norm_eps


class       __lowerCAmelCase					(       _UpperCamelCase      ):


	'''simple docstring'''





	_A													=				version.parse("1.11"      )

	@property
	def    snake_case__(						self:						Optional[int]			):
			return OrderedDict([('pixel_values', {0: 'batch'})]			)

	@property
	def    snake_case__(						self:						Dict			):
			if self.task == "image-classification":
					return OrderedDict([('logits', {0: 'batch'})]			)
			else:
					return OrderedDict([('last_hidden_state', {0: 'batch'}), ('pooler_output', {0: 'batch'})]			)





	@property
	def    snake_case__(						self:						List[Any]			):
			return 1E-4


