


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


    def        A   (  self      :							Tuple):
           _A					:     str     = tempfile.mkdtemp()
           _A					:     str     = 8

           # DPR tok
           _A					:     List[str]     = [
               '[UNK]',
               '[CLS]',
               '[SEP]',
               '[PAD]',
               '[MASK]',
               'want',
               '##want',
               '##ed',
               'wa',
               'un',
               'runn',
               '##ing',
               ',',
               'low',
               'lowest',
           ]
           _A					:     Tuple     = os.path.join(self.tmpdirname				,							'dpr_tokenizer')
           os.makedirs(__UpperCamelCase				,							exist_ok=__UpperCamelCase)
           _A					:     Dict     = 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
           _A					:     Optional[int]     = [
               '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(__UpperCamelCase				,							range(len(__UpperCamelCase))))
           _A					:     str     = ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
           _A					:     Dict     = {'unk_token': '<unk>'}

           _A					:     Dict     = os.path.join(self.tmpdirname				,							'bart_tokenizer')
           os.makedirs(__UpperCamelCase				,							exist_ok=__UpperCamelCase)
           _A					:     int     = os.path.join(__UpperCamelCase				,							BART_VOCAB_FILES_NAMES['vocab_file'])
           _A					:     Dict     = 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        A   (  self      :							Union[str, Any]):
           return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,							'dpr_tokenizer'))


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


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


    def        A   (  self      :							List[str]):
           shutil.rmtree(self.tmpdirname)


    def        A   (  self      :							Optional[int]):
           _A					:     Tuple     = 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        A   (  self      :							int):
           _A					:     Optional[int]     = 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					:     List[Any]     = dataset
                  _A					:     Union[str, Any]     = RagRetriever(
                      __UpperCamelCase				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer()				,							)
           return retriever


    def        A   (  self      :							List[str]				,							SCREAMING_SNAKE_CASE      :							bool):
           _A					:     Tuple     = self.get_dummy_dataset()
           _A					:     str     = RagConfig(
               retrieval_vector_size=self.retrieval_vector_size				,							question_encoder=DPRConfig().to_dict()				,							generator=BartConfig().to_dict()				,							index_name='custom'				,							)
           if from_disk:
                  _A					:     Any     = os.path.join(self.tmpdirname				,							'dataset')
                  _A					:     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					:     Any     = RagRetriever(
                      __UpperCamelCase				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer()				,							)
           else:
                  _A					:     str     = 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        A   (  self      :							Any):
           _A					:     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					:     Union[str, 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					:     Dict     = os.path.join(self.tmpdirname				,							'psgs_w100.tsv.pkl')
           _A					:     int     = {sample['id']: [sample['text'], sample['title']] for sample in dataset}
           pickle.dump(__UpperCamelCase				,							open(__UpperCamelCase				,							'wb'))

           _A					:     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				,							)
           _A					:     int     = RagRetriever(
               __UpperCamelCase				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer())
           return retriever


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


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


    def        A   (  self      :							Dict):
           _A					:     Optional[int]     = 1
           _A					:     Union[str, Any]     = self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCamelCase)
           _A					:     Optional[Any]     = np.array(
               [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]				,							dtype=np.floataa)
           _A  ,			_A  ,			_A					:     Dict     = 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        A   (  self      :							str):
           _A					:     Optional[Any]     = self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCamelCase)
           with tempfile.TemporaryDirectory() as tmp_dirname:
                  retriever.save_pretrained(__UpperCamelCase)
                  _A					:     Union[str, Any]     = RagRetriever.from_pretrained(__UpperCamelCase)
                  self.assertIsInstance(__UpperCamelCase				,							__UpperCamelCase)
                  _A					:     Optional[Any]     = np.array(
                      [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]				,							dtype=np.floataa)
                  _A					:     Union[str, Any]     = retriever.retrieve(__UpperCamelCase				,							n_docs=1)
                  self.assertTrue(out is not None)


    def        A   (  self      :							Optional[Any]):
           _A					:     Optional[int]     = 1
           _A					:     Optional[int]     = self.get_dummy_legacy_index_retriever()
           _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(__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        A   (  self      :							Dict):
           _A					:     Tuple     = self.get_dummy_legacy_index_retriever()
           with tempfile.TemporaryDirectory() as tmp_dirname:
                  retriever.save_pretrained(__UpperCamelCase)
                  _A					:     Union[str, Any]     = RagRetriever.from_pretrained(__UpperCamelCase)
                  self.assertIsInstance(__UpperCamelCase				,							__UpperCamelCase)
                  _A					:     Tuple     = np.array(
                      [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]				,							dtype=np.floataa)
                  _A					:     Optional[int]     = retriever.retrieve(__UpperCamelCase				,							n_docs=1)
                  self.assertTrue(out is not None)


    @require_torch
    @require_tokenizers
    @require_sentencepiece
    def        A   (  self      :							Optional[Any]):
           import torch

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

           _A					:     int     = retriever(
               __UpperCamelCase				,							__UpperCamelCase				,							prefix=retriever.config.generator.prefix				,							n_docs=__UpperCamelCase				,							return_tensors='pt'				,							)
           _A  ,			_A  ,			_A  ,			_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(__UpperCamelCase				,							torch.Tensor)
           self.assertIsInstance(__UpperCamelCase				,							torch.Tensor)
           self.assertIsInstance(__UpperCamelCase				,							torch.Tensor)


    @require_torch
    @require_tokenizers
    @require_sentencepiece
    def        A   (  self      :							List[Any]):
           _A					:     Optional[Any]     = self.get_dpr_ctx_encoder_tokenizer()
           _A					:     List[Any]     = 1
           _A					:     List[Any]     = self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCamelCase)
           retriever.set_ctx_encoder_tokenizer(__UpperCamelCase)

           _A					:     Tuple     = [[5, 7], [10, 11]]
           _A					:     Union[str, Any]     = np.array(
               [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]				,							dtype=np.floataa)
           _A					:     Optional[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.










'''simple docstring'''







from __future__ import annotations


class 			__lowerCamelCase			:
    """simple docstring"""


    def __init__(  self      :							Optional[Any]				,							SCREAMING_SNAKE_CASE      :							Tuple=None):
           _A					:     Any     = data
           _A					:     Optional[Any]     = None


    def __repr__(  self      :							List[str]):
           _A					:     List[Any]     = []
           _A					:     Any     = self
           while temp:
                  string_rep.append(F'{temp.data}')
                  _A					:     List[Any]     = temp.next
           return "->".join(SCREAMING_SNAKE_CASE)






def 				lowerCAmelCase__	(    lowerCamelCase					:      list       ):
       if not elements_list:
              raise Exception('The Elements List is empty'       )

       _A					:     Union[str, Any]     = Node(elements_list[0]       )
       for i in range(1  ,len(lowerCamelCase       )       ):
              _A					:     Dict     = Node(elements_list[i]       )
              _A					:     int     = current.next
       return head






def 				lowerCAmelCase__	(    lowerCamelCase					:      Node       ):
       if head_node is not None and isinstance(lowerCamelCase  ,lowerCamelCase       ):
              print_reverse(head_node.next       )
              print(head_node.data       )






def 				lowerCAmelCase__	(    ):
       from doctest import testmod

       testmod()

       _A					:     List[str]     = make_linked_list([14, 52, 14, 12, 43]       )
       print('Linked List:'       )
       print(lowerCamelCase       )
       print('Elements in Reverse:'       )
       print_reverse(lowerCamelCase       )


if __name__ == "__main__":
   main()






