
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 			UpperCAmelCase_     (					_A	):


  '''simple docstring'''







  def 				_lowercase				(					self     :					Union[str, Any]       )     ->      List[Any]:







    """simple docstring"""






    __magic_name__								=			tempfile.mkdtemp()
    __magic_name__								=			8

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

    __magic_name__								=			os.path.join(self.tmpdirname     ,				"""bart_tokenizer"""       )
    os.makedirs(UpperCamelCase__     ,				exist_ok=UpperCamelCase__       )
    __magic_name__								=			os.path.join(UpperCamelCase__     ,				BART_VOCAB_FILES_NAMES["""vocab_file"""]       )
    __magic_name__								=			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 				_lowercase				(					self     :					List[str]       )     ->      DPRQuestionEncoderTokenizer:







    """simple docstring"""






    return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,				"""dpr_tokenizer"""       )       )
  def 				_lowercase				(					self     :					Tuple       )     ->      DPRContextEncoderTokenizer:







    """simple docstring"""






    return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,				"""dpr_tokenizer"""       )       )
  def 				_lowercase				(					self     :					Any       )     ->      BartTokenizer:







    """simple docstring"""






    return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,				"""bart_tokenizer"""       )       )
  def 				_lowercase				(					self     :					Optional[int]       )     ->      Union[str, Any]:







    """simple docstring"""






    shutil.rmtree(self.tmpdirname       )
  def 				_lowercase				(					self     :					List[Any]       )     ->      Any:







    """simple docstring"""






    __magic_name__								=			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 				_lowercase				(					self     :					Union[str, Any]       )     ->      str:







    """simple docstring"""






    __magic_name__								=			self.get_dummy_dataset()
    __magic_name__								=			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:
      __magic_name__								=			dataset
      __magic_name__								=			RagRetriever(
          UpperCamelCase__     ,				question_encoder_tokenizer=self.get_dpr_tokenizer()     ,				generator_tokenizer=self.get_bart_tokenizer()     ,				)
    return retriever
  def 				_lowercase				(					self     :					List[str]     ,				UpperCamelCase__     :					bool       )     ->      int:







    """simple docstring"""






    __magic_name__								=			self.get_dummy_dataset()
    __magic_name__								=			RagConfig(
        retrieval_vector_size=self.retrieval_vector_size     ,				question_encoder=DPRConfig().to_dict()     ,				generator=BartConfig().to_dict()     ,				index_name="""custom"""     ,				)
    if from_disk:
      __magic_name__								=			os.path.join(self.tmpdirname     ,				"""dataset"""       )
      __magic_name__								=			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
      __magic_name__								=			RagRetriever(
          UpperCamelCase__     ,				question_encoder_tokenizer=self.get_dpr_tokenizer()     ,				generator_tokenizer=self.get_bart_tokenizer()     ,				)
    else:
      __magic_name__								=			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 				_lowercase				(					self     :					List[str]       )     ->      int:







    """simple docstring"""






    __magic_name__								=			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       )

    __magic_name__								=			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"""       )       )

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

    __magic_name__								=			RagConfig(
        retrieval_vector_size=self.retrieval_vector_size     ,				question_encoder=DPRConfig().to_dict()     ,				generator=BartConfig().to_dict()     ,				index_name="""legacy"""     ,				index_path=self.tmpdirname     ,				)
    __magic_name__								=			RagRetriever(
        UpperCamelCase__     ,				question_encoder_tokenizer=self.get_dpr_tokenizer()     ,				generator_tokenizer=self.get_bart_tokenizer()       )
    return retriever
  def 				_lowercase				(					self     :					List[str]       )     ->      str:







    """simple docstring"""






    __magic_name__								=			1
    __magic_name__								=			self.get_dummy_canonical_hf_index_retriever()
    __magic_name__								=			np.array(
        [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]     ,				dtype=np.floataa       )
    __magic_name__	,       __magic_name__	,       __magic_name__								=			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 				_lowercase				(					self     :					Tuple       )     ->      Any:







    """simple docstring"""






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







    """simple docstring"""






    __magic_name__								=			1
    __magic_name__								=			self.get_dummy_custom_hf_index_retriever(from_disk=UpperCamelCase__       )
    __magic_name__								=			np.array(
        [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]     ,				dtype=np.floataa       )
    __magic_name__	,       __magic_name__	,       __magic_name__								=			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 				_lowercase				(					self     :					Any       )     ->      str:







    """simple docstring"""






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







    """simple docstring"""






    __magic_name__								=			1
    __magic_name__								=			self.get_dummy_custom_hf_index_retriever(from_disk=UpperCamelCase__       )
    __magic_name__								=			np.array(
        [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]     ,				dtype=np.floataa       )
    __magic_name__	,       __magic_name__	,       __magic_name__								=			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 				_lowercase				(					self     :					str       )     ->      List[Any]:







    """simple docstring"""






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







    """simple docstring"""






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







    """simple docstring"""






    __magic_name__								=			self.get_dummy_legacy_index_retriever()
    with tempfile.TemporaryDirectory() as tmp_dirname:
      retriever.save_pretrained(UpperCamelCase__       )
      __magic_name__								=			RagRetriever.from_pretrained(UpperCamelCase__       )
      self.assertIsInstance(UpperCamelCase__     ,				UpperCamelCase__       )
      __magic_name__								=			np.array(
          [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]     ,				dtype=np.floataa       )
      __magic_name__								=			retriever.retrieve(UpperCamelCase__     ,				n_docs=1       )
      self.assertTrue(out is not None       )
  @require_torch
  @require_tokenizers
  @require_sentencepiece
  def 				_lowercase				(					self     :					str       )     ->      List[Any]:







    """simple docstring"""






    import torch

    __magic_name__								=			1
    __magic_name__								=			self.get_dummy_canonical_hf_index_retriever()
    __magic_name__								=			[[5, 7], [10, 11]]
    __magic_name__								=			np.array(
        [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]     ,				dtype=np.floataa       )
    __magic_name__								=			retriever(UpperCamelCase__     ,				UpperCamelCase__     ,				prefix=retriever.config.generator.prefix     ,				n_docs=UpperCamelCase__       )
    __magic_name__	,       __magic_name__	,       __magic_name__								=			(
        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       )

    __magic_name__								=			retriever(
        UpperCamelCase__     ,				UpperCamelCase__     ,				prefix=retriever.config.generator.prefix     ,				n_docs=UpperCamelCase__     ,				return_tensors="""pt"""     ,				)
    __magic_name__	,       __magic_name__	,       __magic_name__	,       __magic_name__								=			(  # 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 				_lowercase				(					self     :					int       )     ->      List[Any]:







    """simple docstring"""






    __magic_name__								=			self.get_dpr_ctx_encoder_tokenizer()
    __magic_name__								=			1
    __magic_name__								=			self.get_dummy_custom_hf_index_retriever(from_disk=UpperCamelCase__       )
    retriever.set_ctx_encoder_tokenizer(UpperCamelCase__       )

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



import warnings

from ...utils import logging
from .image_processing_beit import BeitImageProcessor


__lowerCAmelCase     :	int       			=  logging.get_logger(__name__)




class 			UpperCAmelCase_     (					_A	):


  '''simple docstring'''







  def __init__(					self     :					int     ,				*UpperCamelCase__     :					Optional[Any]     ,				**UpperCamelCase__     :					Optional[Any]       )     ->      None:







    """simple docstring"""






    warnings.warn(
        """The class BeitFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"""
        """ use BeitImageProcessor instead."""     ,				UpperCamelCase__     ,				)
    super().__init__(*UpperCamelCase__     ,				**UpperCamelCase__       )

