

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       __snake_case  (    lowerCAmelCase__       ):

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						tempfile.mkdtemp()
      SCREAMING_SNAKE_CASE_									=						8

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

      SCREAMING_SNAKE_CASE_									=						os.path.join(self.tmpdirname						,		'bart_tokenizer')
      os.makedirs(_A						,		exist_ok=_A)
      SCREAMING_SNAKE_CASE_									=						os.path.join(_A						,		BART_VOCAB_FILES_NAMES['vocab_file'])
      SCREAMING_SNAKE_CASE_									=						os.path.join(_A						,		BART_VOCAB_FILES_NAMES['merges_file'])
      with open(self.vocab_file						,		'w'						,		encoding='utf-8') as fp:
          fp.write(json.dumps(_A) + '\n')
      with open(self.merges_file						,		'w'						,		encoding='utf-8') as fp:
          fp.write('\n'.join(_A))

  def 						lowerCAmelCase__     (					self):
      return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname						,		'dpr_tokenizer'))

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

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

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

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						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 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_dataset()
      SCREAMING_SNAKE_CASE_									=						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:
          SCREAMING_SNAKE_CASE_									=						dataset
          SCREAMING_SNAKE_CASE_									=						RagRetriever(
              _A						,		question_encoder_tokenizer=self.get_dpr_tokenizer()						,		generator_tokenizer=self.get_bart_tokenizer()						,		)
      return retriever

  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_dataset()
      SCREAMING_SNAKE_CASE_									=						RagConfig(
          retrieval_vector_size=self.retrieval_vector_size						,		question_encoder=DPRConfig().to_dict()						,		generator=BartConfig().to_dict()						,		index_name='custom'						,		)
      if from_disk:
          SCREAMING_SNAKE_CASE_									=						os.path.join(self.tmpdirname						,		'dataset')
          SCREAMING_SNAKE_CASE_									=						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
          SCREAMING_SNAKE_CASE_									=						RagRetriever(
              _A						,		question_encoder_tokenizer=self.get_dpr_tokenizer()						,		generator_tokenizer=self.get_bart_tokenizer()						,		)
      else:
          SCREAMING_SNAKE_CASE_									=						RagRetriever(
              _A						,		question_encoder_tokenizer=self.get_dpr_tokenizer()						,		generator_tokenizer=self.get_bart_tokenizer()						,		index=CustomHFIndex(config.retrieval_vector_size						,		_A)						,		)
      return retriever

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						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)

      SCREAMING_SNAKE_CASE_									=						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'))

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

      SCREAMING_SNAKE_CASE_									=						RagConfig(
          retrieval_vector_size=self.retrieval_vector_size						,		question_encoder=DPRConfig().to_dict()						,		generator=BartConfig().to_dict()						,		index_name='legacy'						,		index_path=self.tmpdirname						,		)
      SCREAMING_SNAKE_CASE_									=						RagRetriever(
          _A						,		question_encoder_tokenizer=self.get_dpr_tokenizer()						,		generator_tokenizer=self.get_bart_tokenizer())
      return retriever

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						1
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_canonical_hf_index_retriever()
      SCREAMING_SNAKE_CASE_									=						np.array(
          [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=_A)
      self.assertEqual(retrieved_doc_embeds.shape						,		(2, n_docs, self.retrieval_vector_size))
      self.assertEqual(len(_A)						,		2)
      self.assertEqual(sorted(doc_dicts[0])						,		['embeddings', 'id', 'text', 'title'])
      self.assertEqual(len(doc_dicts[0]['id'])						,		_A)
      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 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						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:
              SCREAMING_SNAKE_CASE_									=						self.get_dummy_dataset()
              retriever.save_pretrained(_A)
              SCREAMING_SNAKE_CASE_									=						RagRetriever.from_pretrained(_A)
              self.assertIsInstance(_A						,		_A)
              SCREAMING_SNAKE_CASE_									=						np.array(
                  [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
              SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=1)
              self.assertTrue(out is not None)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						1
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_custom_hf_index_retriever(from_disk=_A)
      SCREAMING_SNAKE_CASE_									=						np.array(
          [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=_A)
      self.assertEqual(retrieved_doc_embeds.shape						,		(2, n_docs, self.retrieval_vector_size))
      self.assertEqual(len(_A)						,		2)
      self.assertEqual(sorted(doc_dicts[0])						,		['embeddings', 'id', 'text', 'title'])
      self.assertEqual(len(doc_dicts[0]['id'])						,		_A)
      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 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_custom_hf_index_retriever(from_disk=_A)
      with tempfile.TemporaryDirectory() as tmp_dirname:
          retriever.save_pretrained(_A)
          SCREAMING_SNAKE_CASE_									=						RagRetriever.from_pretrained(_A)
          self.assertIsInstance(_A						,		_A)
          SCREAMING_SNAKE_CASE_									=						np.array(
              [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
          SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=1)
          self.assertTrue(out is not None)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						1
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_custom_hf_index_retriever(from_disk=_A)
      SCREAMING_SNAKE_CASE_									=						np.array(
          [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=_A)
      self.assertEqual(retrieved_doc_embeds.shape						,		(2, n_docs, self.retrieval_vector_size))
      self.assertEqual(len(_A)						,		2)
      self.assertEqual(sorted(doc_dicts[0])						,		['embeddings', 'id', 'text', 'title'])
      self.assertEqual(len(doc_dicts[0]['id'])						,		_A)
      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 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_custom_hf_index_retriever(from_disk=_A)
      with tempfile.TemporaryDirectory() as tmp_dirname:
          retriever.save_pretrained(_A)
          SCREAMING_SNAKE_CASE_									=						RagRetriever.from_pretrained(_A)
          self.assertIsInstance(_A						,		_A)
          SCREAMING_SNAKE_CASE_									=						np.array(
              [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
          SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=1)
          self.assertTrue(out is not None)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						1
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_legacy_index_retriever()
      SCREAMING_SNAKE_CASE_									=						np.array(
          [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=_A)
      self.assertEqual(retrieved_doc_embeds.shape						,		(2, n_docs, self.retrieval_vector_size))
      self.assertEqual(len(_A)						,		2)
      self.assertEqual(sorted(doc_dicts[0])						,		['text', 'title'])
      self.assertEqual(len(doc_dicts[0]['text'])						,		_A)
      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 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_legacy_index_retriever()
      with tempfile.TemporaryDirectory() as tmp_dirname:
          retriever.save_pretrained(_A)
          SCREAMING_SNAKE_CASE_									=						RagRetriever.from_pretrained(_A)
          self.assertIsInstance(_A						,		_A)
          SCREAMING_SNAKE_CASE_									=						np.array(
              [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
          SCREAMING_SNAKE_CASE_									=						retriever.retrieve(_A						,		n_docs=1)
          self.assertTrue(out is not None)

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

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

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


  @require_torch
  @require_tokenizers
  @require_sentencepiece
  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.get_dpr_ctx_encoder_tokenizer()
      SCREAMING_SNAKE_CASE_									=						1
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_custom_hf_index_retriever(from_disk=_A)
      retriever.set_ctx_encoder_tokenizer(_A)

      SCREAMING_SNAKE_CASE_									=						[[5, 7], [10, 11]]
      SCREAMING_SNAKE_CASE_									=						np.array(
          [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]						,		dtype=np.floataa)
      SCREAMING_SNAKE_CASE_									=						retriever(_A						,		_A						,		prefix=retriever.config.generator.prefix						,		n_docs=_A)

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





from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__  :      str			=    logging.get_logger(__name__)

UpperCamelCase__  :      Optional[int]			=    {
    "facebook/dpr-ctx_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-single-nq-base": (
        "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-ctx_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-multiset-base": (
        "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"
    ),
}



class       __snake_case  (    lowerCAmelCase__       ):
  __lowerCAmelCase     :	Optional[int]						 =  'dpr'

  def __init__(					self						,		_A=30522						,		_A=768						,		_A=12						,		_A=12						,		_A=3072						,		_A="gelu"						,		_A=0.1						,		_A=0.1						,		_A=512						,		_A=2						,		_A=0.0_2						,		_A=1E-12						,		_A=0						,		_A="absolute"						,		_A = 0						,		**_A						,		):
      super().__init__(pad_token_id=_A						,		**_A)

      SCREAMING_SNAKE_CASE_									=						vocab_size
      SCREAMING_SNAKE_CASE_									=						hidden_size
      SCREAMING_SNAKE_CASE_									=						num_hidden_layers
      SCREAMING_SNAKE_CASE_									=						num_attention_heads
      SCREAMING_SNAKE_CASE_									=						hidden_act
      SCREAMING_SNAKE_CASE_									=						intermediate_size
      SCREAMING_SNAKE_CASE_									=						hidden_dropout_prob
      SCREAMING_SNAKE_CASE_									=						attention_probs_dropout_prob
      SCREAMING_SNAKE_CASE_									=						max_position_embeddings
      SCREAMING_SNAKE_CASE_									=						type_vocab_size
      SCREAMING_SNAKE_CASE_									=						initializer_range
      SCREAMING_SNAKE_CASE_									=						layer_norm_eps
      SCREAMING_SNAKE_CASE_									=						projection_dim
      SCREAMING_SNAKE_CASE_									=						position_embedding_type


