





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__  ):





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						Dict      		=  tempfile.mkdtemp()
      a__	:						Optional[Any]      		=  8

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

      a__	:						int      		=  os.path.join(self.tmpdirname					,			"bart_tokenizer"      )
      os.makedirs(snake_case__					,			exist_ok=snake_case__      )
      a__	:						List[str]      		=  os.path.join(snake_case__					,			BART_VOCAB_FILES_NAMES["vocab_file"]      )
      a__	:						str      		=  os.path.join(snake_case__					,			BART_VOCAB_FILES_NAMES["merges_file"]      )
      with open(self.vocab_file					,			"w"					,			encoding="utf-8"      ) as fp:
         fp.write(json.dumps(snake_case__      ) + "\n"      )
      with open(self.merges_file					,			"w"					,			encoding="utf-8"      ) as fp:
         fp.write("\n".join(snake_case__      )      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      shutil.rmtree(self.tmpdirname      )





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Dict      		=  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 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





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





   def 	_snake_case ( self					,			snake_case      )				->  Union[str, Any]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      a__	:						List[str]      		=  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      )

      a__	:						Optional[int]      		=  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"      )      )

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

      a__	:						List[Any]      		=  RagConfig(
          retrieval_vector_size=self.retrieval_vector_size					,			question_encoder=DPRConfig().to_dict()					,			generator=BartConfig().to_dict()					,			index_name="legacy"					,			index_path=self.tmpdirname					,			)
      a__	:						Optional[Any]      		=  RagRetriever(
          snake_case__					,			question_encoder_tokenizer=self.get_dpr_tokenizer()					,			generator_tokenizer=self.get_bart_tokenizer()      )
      return retriever





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      a__	:						Dict      		=  1
      a__	:						Optional[Any]      		=  self.get_dummy_canonical_hf_index_retriever()
      a__	:						Union[str, Any]      		=  np.array(
          [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					,			dtype=np.floataa      )
      a__	:						Optional[int]      		=  retriever.retrieve(snake_case__					,			n_docs=snake_case__      )
      self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)      )
      self.assertEqual(len(snake_case__      )					,			2      )
      self.assertEqual(sorted(doc_dicts[0]      )					,			["embeddings", "id", "text", "title"]      )
      self.assertEqual(len(doc_dicts[0]["id"]      )					,			snake_case__      )
      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 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						List[Any]      		=  1
      a__	:						Any      		=  self.get_dummy_custom_hf_index_retriever(from_disk=snake_case__      )
      a__	:						Union[str, Any]      		=  np.array(
          [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					,			dtype=np.floataa      )
      a__	:						Any      		=  retriever.retrieve(snake_case__					,			n_docs=snake_case__      )
      self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)      )
      self.assertEqual(len(snake_case__      )					,			2      )
      self.assertEqual(sorted(doc_dicts[0]      )					,			["embeddings", "id", "text", "title"]      )
      self.assertEqual(len(doc_dicts[0]["id"]      )					,			snake_case__      )
      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 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						Any      		=  1
      a__	:						str      		=  self.get_dummy_custom_hf_index_retriever(from_disk=snake_case__      )
      a__	:						List[str]      		=  np.array(
          [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					,			dtype=np.floataa      )
      a__	:						Optional[int]      		=  retriever.retrieve(snake_case__					,			n_docs=snake_case__      )
      self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)      )
      self.assertEqual(len(snake_case__      )					,			2      )
      self.assertEqual(sorted(doc_dicts[0]      )					,			["embeddings", "id", "text", "title"]      )
      self.assertEqual(len(doc_dicts[0]["id"]      )					,			snake_case__      )
      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 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      a__	:						Dict      		=  1
      a__	:						int      		=  self.get_dummy_legacy_index_retriever()
      a__	:						Optional[Any]      		=  np.array(
          [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					,			dtype=np.floataa      )
      a__	:						Optional[Any]      		=  retriever.retrieve(snake_case__					,			n_docs=snake_case__      )
      self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)      )
      self.assertEqual(len(snake_case__      )					,			2      )
      self.assertEqual(sorted(doc_dicts[0]      )					,			["text", "title"]      )
      self.assertEqual(len(doc_dicts[0]["text"]      )					,			snake_case__      )
      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 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





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





   @require_torch
   @require_tokenizers
   @require_sentencepiece
   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      import torch

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

      a__	:						Any      		=  retriever(
          snake_case__					,			snake_case__					,			prefix=retriever.config.generator.prefix					,			n_docs=snake_case__					,			return_tensors="pt"					,			)
      a__	:						Tuple      		=  (  # 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(snake_case__					,			torch.Tensor      )
      self.assertIsInstance(snake_case__					,			torch.Tensor      )
      self.assertIsInstance(snake_case__					,			torch.Tensor      )





   @require_torch
   @require_tokenizers
   @require_sentencepiece
   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      a__	:						List[str]      		=  self.get_dpr_ctx_encoder_tokenizer()
      a__	:						Dict      		=  1
      a__	:						Any      		=  self.get_dummy_custom_hf_index_retriever(from_disk=snake_case__      )
      retriever.set_ctx_encoder_tokenizer(snake_case__      )

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

      self.assertEqual(
          len(snake_case__      )					,			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")      )					,			snake_case__      )  # 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


SCREAMING_SNAKE_CASE__       :			Dict			      =						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__       :			Any			      =						{
    """google/mobilenet_v2_1.4_224""": """https://huggingface.co/google/mobilenet_v2_1.4_224/resolve/main/config.json""",
    """google/mobilenet_v2_1.0_224""": """https://huggingface.co/google/mobilenet_v2_1.0_224/resolve/main/config.json""",
    """google/mobilenet_v2_0.75_160""": """https://huggingface.co/google/mobilenet_v2_0.75_160/resolve/main/config.json""",
    """google/mobilenet_v2_0.35_96""": """https://huggingface.co/google/mobilenet_v2_0.35_96/resolve/main/config.json""",
    # See all MobileNetV2 models at https://huggingface.co/models?filter=mobilenet_v2
}







class   __lowerCAmelCase			(					_UpperCamelCase  ):
   _UpperCamelCase :						Optional[int]     		=				"""mobilenet_v2"""





   def __init__( self					,			snake_case=3					,			snake_case=224					,			snake_case=1.0					,			snake_case=8					,			snake_case=8					,			snake_case=6					,			snake_case=32					,			snake_case=True					,			snake_case=True					,			snake_case="relu6"					,			snake_case=True					,			snake_case=0.8					,			snake_case=0.02					,			snake_case=0.001					,			snake_case=255					,			**snake_case					,			)				->  int:
      """simple docstring"""





      super().__init__(**snake_case      )

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

      a__	:						str      		=  num_channels
      a__	:						Dict      		=  image_size
      a__	:						Any      		=  depth_multiplier
      a__	:						str      		=  depth_divisible_by
      a__	:						Optional[int]      		=  min_depth
      a__	:						Dict      		=  expand_ratio
      a__	:						str      		=  output_stride
      a__	:						Optional[int]      		=  first_layer_is_expansion
      a__	:						Union[str, Any]      		=  finegrained_output
      a__	:						Union[str, Any]      		=  hidden_act
      a__	:						str      		=  tf_padding
      a__	:						List[Any]      		=  classifier_dropout_prob
      a__	:						List[Any]      		=  initializer_range
      a__	:						Optional[Any]      		=  layer_norm_eps
      a__	:						str      		=  semantic_loss_ignore_index







class   __lowerCAmelCase			(					_UpperCamelCase  ):
   _UpperCamelCase :						Any     		=				version.parse("""1.11"""  )





   @property
   def 	_snake_case ( self      )				->  Mapping[str, Mapping[int, str]]:
      """simple docstring"""





      return OrderedDict([("pixel_values", {0: "batch"})]      )





   @property
   def 	_snake_case ( self      )				->  Mapping[str, Mapping[int, str]]:
      """simple docstring"""





      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      )				->  float:
      """simple docstring"""





      return 1E-4

