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_						(	snake_case_						):







   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    UpperCamelCase          =	tempfile.mkdtemp()
    UpperCamelCase          =	8

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

    UpperCamelCase          =	os.path.join(self.tmpdirname					,					"""bart_tokenizer"""							)
    os.makedirs(lowerCamelCase_					,					exist_ok=lowerCamelCase_							)
    UpperCamelCase          =	os.path.join(lowerCamelCase_					,					BART_VOCAB_FILES_NAMES["""vocab_file"""]							)
    UpperCamelCase          =	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 		lowerCamelCase_     (			self		:		Union[str, Any]							):





    """simple docstring"""


    return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,					"""dpr_tokenizer"""							)							)







   def 		lowerCamelCase_     (			self		:		List[str]							):





    """simple docstring"""


    return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,					"""dpr_tokenizer"""							)							)







   def 		lowerCamelCase_     (			self		:		Union[str, Any]							):





    """simple docstring"""


    return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname					,					"""bart_tokenizer"""							)							)







   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    shutil.rmtree(self.tmpdirname							)







   def 		lowerCamelCase_     (			self		:		Tuple							):





    """simple docstring"""


    UpperCamelCase          =	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		:		Union[str, Any]							):





    """simple docstring"""


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







   def 		lowerCamelCase_     (			self		:		Tuple					,					lowerCamelCase_		:		bool							):





    """simple docstring"""


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





    """simple docstring"""


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

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

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

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







   def 		lowerCamelCase_     (			self		:		Dict							):





    """simple docstring"""


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





    """simple docstring"""


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







   def 		lowerCamelCase_     (			self		:		Optional[int]							):





    """simple docstring"""


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





    """simple docstring"""


    UpperCamelCase          =	self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase_							)
    with tempfile.TemporaryDirectory() as tmp_dirname:
     retriever.save_pretrained(lowerCamelCase_							)
     UpperCamelCase          =	RagRetriever.from_pretrained(lowerCamelCase_							)
     self.assertIsInstance(lowerCamelCase_					,					lowerCamelCase_							)
     UpperCamelCase          =	np.array(
         [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]					,					dtype=np.floataa							)
     UpperCamelCase          =	retriever.retrieve(lowerCamelCase_					,					n_docs=1							)
     self.assertTrue(out is not None							)







   def 		lowerCamelCase_     (			self		:		List[Any]							):





    """simple docstring"""


    UpperCamelCase          =	1
    UpperCamelCase          =	self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase_							)
    UpperCamelCase          =	np.array(
        [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]					,					dtype=np.floataa							)
    UpperCamelCase          =	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 		lowerCamelCase_     (			self		:		int							):





    """simple docstring"""


    UpperCamelCase          =	self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase_							)
    with tempfile.TemporaryDirectory() as tmp_dirname:
     retriever.save_pretrained(lowerCamelCase_							)
     UpperCamelCase          =	RagRetriever.from_pretrained(lowerCamelCase_							)
     self.assertIsInstance(lowerCamelCase_					,					lowerCamelCase_							)
     UpperCamelCase          =	np.array(
         [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]					,					dtype=np.floataa							)
     UpperCamelCase          =	retriever.retrieve(lowerCamelCase_					,					n_docs=1							)
     self.assertTrue(out is not None							)







   def 		lowerCamelCase_     (			self		:		Optional[Any]							):





    """simple docstring"""


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





    """simple docstring"""


    UpperCamelCase          =	self.get_dummy_legacy_index_retriever()
    with tempfile.TemporaryDirectory() as tmp_dirname:
     retriever.save_pretrained(lowerCamelCase_							)
     UpperCamelCase          =	RagRetriever.from_pretrained(lowerCamelCase_							)
     self.assertIsInstance(lowerCamelCase_					,					lowerCamelCase_							)
     UpperCamelCase          =	np.array(
         [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]					,					dtype=np.floataa							)
     UpperCamelCase          =	retriever.retrieve(lowerCamelCase_					,					n_docs=1							)
     self.assertTrue(out is not None							)







   @require_torch
   @require_tokenizers
   @require_sentencepiece
   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    import torch

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

    UpperCamelCase          =	retriever(
        lowerCamelCase_					,					lowerCamelCase_					,					prefix=retriever.config.generator.prefix					,					n_docs=lowerCamelCase_					,					return_tensors="""pt"""					,					)
    UpperCamelCase          =	(  # 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 		lowerCamelCase_     (			self		:		Dict							):





    """simple docstring"""


    UpperCamelCase          =	self.get_dpr_ctx_encoder_tokenizer()
    UpperCamelCase          =	1
    UpperCamelCase          =	self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase_							)
    retriever.set_ctx_encoder_tokenizer(lowerCamelCase_							)

    UpperCamelCase          =	[[5, 7], [10, 11]]
    UpperCamelCase          =	np.array(
        [np.ones(self.retrieval_vector_size							), -np.ones(self.retrieval_vector_size							)]					,					dtype=np.floataa							)
    UpperCamelCase          =	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 ...configuration_utils import PretrainedConfig
from ...utils import logging


a_    					=       logging.get_logger(__name__)

a_    					=       {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}





class 		_lowercase						(					snake_case_       ):
 lowercase				     =   'megatron-bert'

 def __init__( self	:							List[str]   ,      snake_case	:							Tuple=2_9_0_5_6   ,      snake_case	:							Dict=1_0_2_4   ,      snake_case	:							Dict=2_4   ,      snake_case	:							Union[str, Any]=1_6   ,      snake_case	:							Optional[int]=4_0_9_6   ,      snake_case	:							Optional[int]="gelu"   ,      snake_case	:							Any=0.1   ,      snake_case	:							Tuple=0.1   ,      snake_case	:							Optional[int]=5_1_2   ,      snake_case	:							List[Any]=2   ,      snake_case	:							Tuple=0.02   ,      snake_case	:							Optional[Any]=1e-12   ,      snake_case	:							str=0   ,      snake_case	:							Optional[int]="absolute"   ,      snake_case	:							Union[str, Any]=True   ,      **snake_case	:							Any   ,      )   ->						List[Any]:

  """simple docstring"""





  super().__init__(pad_token_id=snake_case   ,      **snake_case			)

  UpperCamelCase_					:						Optional[Any] =					vocab_size
  UpperCamelCase_					:						Any =					hidden_size
  UpperCamelCase_					:						Union[str, Any] =					num_hidden_layers
  UpperCamelCase_					:						List[Any] =					num_attention_heads
  UpperCamelCase_					:						str =					hidden_act
  UpperCamelCase_					:						List[str] =					intermediate_size
  UpperCamelCase_					:						List[Any] =					hidden_dropout_prob
  UpperCamelCase_					:						Any =					attention_probs_dropout_prob
  UpperCamelCase_					:						Tuple =					max_position_embeddings
  UpperCamelCase_					:						Dict =					type_vocab_size
  UpperCamelCase_					:						Optional[int] =					initializer_range
  UpperCamelCase_					:						Optional[Any] =					layer_norm_eps
  UpperCamelCase_					:						Dict =					position_embedding_type
  UpperCamelCase_					:						List[str] =					use_cache



