


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   a__	(						lowerCAmelCase_    ):

       def 						UpperCAmelCase(   self      :     Optional[Any]    ):
            a_    :     Union[str, Any]										=    tempfile.mkdtemp()
            a_    :     List[str]										=    8

            # DPR tok
            a_    :     Optional[int]										=    [
                """[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(lowerCamelCase_      ,							exist_ok=lowerCamelCase_    )
            a_    :     Optional[Any]										=    os.path.join(lowerCamelCase_      ,							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_    :     Any										=    [
                """l""",
                """o""",
                """w""",
                """e""",
                """r""",
                """s""",
                """t""",
                """i""",
                """d""",
                """n""",
                """\u0120""",
                """\u0120l""",
                """\u0120n""",
                """\u0120lo""",
                """\u0120low""",
                """er""",
                """\u0120lowest""",
                """\u0120newer""",
                """\u0120wider""",
                """<unk>""",
            ]
            a_    :     Optional[int]										=    dict(zip(lowerCamelCase_      ,							range(len(lowerCamelCase_    )    )    )    )
            a_    :     Tuple										=    ["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
            a_    :     Dict										=    {"""unk_token""": """<unk>"""}

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

       def 						UpperCAmelCase(   self      :     Optional[int]    ):
            return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname      ,							"""dpr_tokenizer"""    )    )

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

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

       def 						UpperCAmelCase(   self      :     str    ):
            shutil.rmtree(self.tmpdirname    )

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

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

       def 						UpperCAmelCase(   self      :     Optional[Any]    ):
            a_    :     Optional[int]										=    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_    :     List[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"""    )    )

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

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

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

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

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

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

       @require_torch
       @require_tokenizers
       @require_sentencepiece
       def 						UpperCAmelCase(   self      :     str    ):
            import torch

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

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

       @require_torch
       @require_tokenizers
       @require_sentencepiece
       def 						UpperCAmelCase(   self      :     Dict    ):
            a_    :     Optional[int]										=    self.get_dpr_ctx_encoder_tokenizer()
            a_    :     Any										=    1
            a_    :     Optional[Any]										=    self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase_    )
            retriever.set_ctx_encoder_tokenizer(lowerCamelCase_    )

            a_    :     int										=    [[5, 7], [1_0, 1_1]]
            a_    :     int										=    np.array(
                [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]      ,							dtype=np.floataa    )
            a_    :     str										=    retriever(lowerCamelCase_      ,							lowerCamelCase_      ,							prefix=retriever.config.generator.prefix      ,							n_docs=lowerCamelCase_    )

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







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__lowerCamelCase        =     {
    '''configuration_pegasus_x''': ['''PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''PegasusXConfig'''],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __lowerCamelCase        =     [
           '''PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''PegasusXForConditionalGeneration''',
           '''PegasusXModel''',
           '''PegasusXPreTrainedModel''',
       ]


if TYPE_CHECKING:
       from .configuration_pegasus_x import PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusXConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_pegasus_x import (
                  PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST,
                  PegasusXForConditionalGeneration,
                  PegasusXModel,
                  PegasusXPreTrainedModel,
              )


else:
       import sys

       __lowerCamelCase        =     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



