





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



 def 				UpperCamelCase			(							self ):
        A__											=						tempfile.mkdtemp()
        A__											=						8

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

        A__											=						os.path.join(self.tmpdirname,'''bart_tokenizer''' )
        os.makedirs(__lowerCamelCase,exist_ok=__lowerCamelCase )
        A__											=						os.path.join(__lowerCamelCase,BART_VOCAB_FILES_NAMES['''vocab_file'''] )
        A__											=						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 ):
        return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname,'''dpr_tokenizer''' ) )



 def 				UpperCamelCase			(							self ):
        return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname,'''dpr_tokenizer''' ) )



 def 				UpperCamelCase			(							self ):
        return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname,'''bart_tokenizer''' ) )



 def 				UpperCamelCase			(							self ):
        shutil.rmtree(self.tmpdirname )



 def 				UpperCamelCase			(							self ):
        A__											=						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 ):
        A__											=						self.get_dummy_dataset()
        A__											=						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__											=						dataset
               A__											=						RagRetriever(
                   __lowerCamelCase,question_encoder_tokenizer=self.get_dpr_tokenizer(),generator_tokenizer=self.get_bart_tokenizer(),)
        return retriever



 def 				UpperCamelCase			(							self,__lowerCamelCase ):
        A__											=						self.get_dummy_dataset()
        A__											=						RagConfig(
            retrieval_vector_size=self.retrieval_vector_size,question_encoder=DPRConfig().to_dict(),generator=BartConfig().to_dict(),index_name='''custom''',)
        if from_disk:
               A__											=						os.path.join(self.tmpdirname,'''dataset''' )
               A__											=						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__											=						RagRetriever(
                   __lowerCamelCase,question_encoder_tokenizer=self.get_dpr_tokenizer(),generator_tokenizer=self.get_bart_tokenizer(),)
        else:
               A__											=						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 ):
        A__											=						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__											=						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__											=						os.path.join(self.tmpdirname,'''psgs_w100.tsv.pkl''' )
        A__											=						{sample['''id''']: [sample['''text'''], sample['''title''']] for sample in dataset}
        pickle.dump(__lowerCamelCase,open(__lowerCamelCase,'''wb''' ) )

        A__											=						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__											=						RagRetriever(
            __lowerCamelCase,question_encoder_tokenizer=self.get_dpr_tokenizer(),generator_tokenizer=self.get_bart_tokenizer() )
        return retriever



 def 				UpperCamelCase			(							self ):
        A__											=						1
        A__											=						self.get_dummy_canonical_hf_index_retriever()
        A__											=						np.array(
            [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
        A__					,	A__					,	A__											=						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 ):
        A__											=						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__											=						self.get_dummy_dataset()
                      retriever.save_pretrained(__lowerCamelCase )
                      A__											=						RagRetriever.from_pretrained(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )
                      A__											=						np.array(
                          [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
                      A__											=						retriever.retrieve(__lowerCamelCase,n_docs=1 )
                      self.assertTrue(out is not None )



 def 				UpperCamelCase			(							self ):
        A__											=						1
        A__											=						self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase )
        A__											=						np.array(
            [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
        A__					,	A__					,	A__											=						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 ):
        A__											=						self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase )
        with tempfile.TemporaryDirectory() as tmp_dirname:
               retriever.save_pretrained(__lowerCamelCase )
               A__											=						RagRetriever.from_pretrained(__lowerCamelCase )
               self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )
               A__											=						np.array(
                   [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
               A__											=						retriever.retrieve(__lowerCamelCase,n_docs=1 )
               self.assertTrue(out is not None )



 def 				UpperCamelCase			(							self ):
        A__											=						1
        A__											=						self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase )
        A__											=						np.array(
            [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
        A__					,	A__					,	A__											=						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 ):
        A__											=						self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase )
        with tempfile.TemporaryDirectory() as tmp_dirname:
               retriever.save_pretrained(__lowerCamelCase )
               A__											=						RagRetriever.from_pretrained(__lowerCamelCase )
               self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )
               A__											=						np.array(
                   [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
               A__											=						retriever.retrieve(__lowerCamelCase,n_docs=1 )
               self.assertTrue(out is not None )



 def 				UpperCamelCase			(							self ):
        A__											=						1
        A__											=						self.get_dummy_legacy_index_retriever()
        A__											=						np.array(
            [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
        A__					,	A__					,	A__											=						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 ):
        A__											=						self.get_dummy_legacy_index_retriever()
        with tempfile.TemporaryDirectory() as tmp_dirname:
               retriever.save_pretrained(__lowerCamelCase )
               A__											=						RagRetriever.from_pretrained(__lowerCamelCase )
               self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )
               A__											=						np.array(
                   [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
               A__											=						retriever.retrieve(__lowerCamelCase,n_docs=1 )
               self.assertTrue(out is not None )



 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def 				UpperCamelCase			(							self ):
        import torch

        A__											=						1
        A__											=						self.get_dummy_canonical_hf_index_retriever()
        A__											=						[[5, 7], [10, 11]]
        A__											=						np.array(
            [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )],dtype=np.floataa )
        A__											=						retriever(__lowerCamelCase,__lowerCamelCase,prefix=retriever.config.generator.prefix,n_docs=__lowerCamelCase )
        A__					,	A__					,	A__											=						(
            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__											=						retriever(
            __lowerCamelCase,__lowerCamelCase,prefix=retriever.config.generator.prefix,n_docs=__lowerCamelCase,return_tensors='''pt''',)
        A__					,	A__					,	A__					,	A__											=						(  # 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 ):
        A__											=						self.get_dpr_ctx_encoder_tokenizer()
        A__											=						1
        A__											=						self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase )
        retriever.set_ctx_encoder_tokenizer(__lowerCamelCase )

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












import numpy as np


def 	UpperCamelCase__(       UpperCamelCase__			:       np.array  )->np.array:
       return (2 / (1 + np.exp(-2 * vector  ))) - 1


if __name__ == "__main__":
       import doctest

       doctest.testmod()





