


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 	UpperCAmelCase_					(					UpperCamelCase    ):

   '''simple docstring'''




   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Any    =				tempfile.mkdtemp()
     lowerCamelCase :					Union[str, Any]    =				8

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

     lowerCamelCase :					Tuple    =				os.path.join(self.tmpdirname						,				"bart_tokenizer" )
     os.makedirs(__A						,				exist_ok=__A )
     lowerCamelCase :					Optional[Any]    =				os.path.join(__A						,				BART_VOCAB_FILES_NAMES["vocab_file"] )
     lowerCamelCase :					Dict    =				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 _snake_case    (		self ):






     """simple docstring"""





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






   def _snake_case    (		self ):






     """simple docstring"""





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






   def _snake_case    (		self ):






     """simple docstring"""





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






   def _snake_case    (		self ):






     """simple docstring"""





     shutil.rmtree(self.tmpdirname )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					List[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 _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Optional[Any]    =				self.get_dummy_dataset()
     lowerCamelCase :					Dict    =				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:
       lowerCamelCase :					Optional[int]    =				dataset
       lowerCamelCase :					List[Any]    =				RagRetriever(
           __A						,				question_encoder_tokenizer=self.get_dpr_tokenizer()						,				generator_tokenizer=self.get_bart_tokenizer()						,				)
     return retriever






   def _snake_case    (		self						,				__A ):






     """simple docstring"""





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






     """simple docstring"""





     lowerCamelCase :					List[Any]    =				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 )

     lowerCamelCase :					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" ) )

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

     lowerCamelCase :					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						,				)
     lowerCamelCase :					Tuple    =				RagRetriever(
         __A						,				question_encoder_tokenizer=self.get_dpr_tokenizer()						,				generator_tokenizer=self.get_bart_tokenizer() )
     return retriever






   def _snake_case    (		self ):






     """simple docstring"""





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






     """simple docstring"""





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






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Union[str, Any]    =				1
     lowerCamelCase :					Any    =				self.get_dummy_custom_hf_index_retriever(from_disk=__A )
     lowerCamelCase :					int    =				np.array(
         [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]						,				dtype=np.floataa )
     lowerCamelCase    ,	lowerCamelCase    ,	lowerCamelCase :					List[str]    =				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 _snake_case    (		self ):






     """simple docstring"""





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






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Optional[int]    =				1
     lowerCamelCase :					str    =				self.get_dummy_custom_hf_index_retriever(from_disk=__A )
     lowerCamelCase :					Dict    =				np.array(
         [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]						,				dtype=np.floataa )
     lowerCamelCase    ,	lowerCamelCase    ,	lowerCamelCase :					str    =				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 _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Optional[Any]    =				self.get_dummy_custom_hf_index_retriever(from_disk=__A )
     with tempfile.TemporaryDirectory() as tmp_dirname:
       retriever.save_pretrained(__A )
       lowerCamelCase :					List[str]    =				RagRetriever.from_pretrained(__A )
       self.assertIsInstance(__A						,				__A )
       lowerCamelCase :					Union[str, Any]    =				np.array(
           [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]						,				dtype=np.floataa )
       lowerCamelCase :					int    =				retriever.retrieve(__A						,				n_docs=1 )
       self.assertTrue(out is not None )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Dict    =				1
     lowerCamelCase :					str    =				self.get_dummy_legacy_index_retriever()
     lowerCamelCase :					Tuple    =				np.array(
         [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]						,				dtype=np.floataa )
     lowerCamelCase    ,	lowerCamelCase    ,	lowerCamelCase :					Tuple    =				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 _snake_case    (		self ):






     """simple docstring"""





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






   @require_torch
   @require_tokenizers
   @require_sentencepiece
   def _snake_case    (		self ):






     """simple docstring"""





     import torch

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

     lowerCamelCase :					Optional[Any]    =				retriever(
         __A						,				__A						,				prefix=retriever.config.generator.prefix						,				n_docs=__A						,				return_tensors="pt"						,				)
     lowerCamelCase    ,	lowerCamelCase    ,	lowerCamelCase    ,	lowerCamelCase :					Optional[int]    =				(  # 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 _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					int    =				self.get_dpr_ctx_encoder_tokenizer()
     lowerCamelCase :					Tuple    =				1
     lowerCamelCase :					List[Any]    =				self.get_dummy_custom_hf_index_retriever(from_disk=__A )
     retriever.set_ctx_encoder_tokenizer(__A )

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







import json
import os
import unittest

from transformers.models.biogpt.tokenization_biogpt import VOCAB_FILES_NAMES, BioGptTokenizer
from transformers.testing_utils import slow

from ...test_tokenization_common import TokenizerTesterMixin
class 	UpperCAmelCase_					(					UpperCamelCase		, unittest.TestCase    ):

   '''simple docstring'''




   __A :      List[Any]     		=		BioGptTokenizer
   __A :      Optional[int]     		=		False






   def _snake_case    (		self ):






     """simple docstring"""





     super().setUp()

     # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
     lowerCamelCase :					Union[str, Any]    =				[
         "l",
         "o",
         "w",
         "e",
         "r",
         "s",
         "t",
         "i",
         "d",
         "n",
         "w</w>",
         "r</w>",
         "t</w>",
         "lo",
         "low",
         "er</w>",
         "low</w>",
         "lowest</w>",
         "newer</w>",
         "wider</w>",
         "<unk>",
     ]
     lowerCamelCase :					str    =				dict(zip(__A						,				range(len(__A ) ) ) )
     lowerCamelCase :					Dict    =				["l o 123", "lo w 1456", "e r</w> 1789", ""]

     lowerCamelCase :					Dict    =				os.path.join(self.tmpdirname						,				VOCAB_FILES_NAMES["vocab_file"] )
     lowerCamelCase :					int    =				os.path.join(self.tmpdirname						,				VOCAB_FILES_NAMES["merges_file"] )
     with open(self.vocab_file						,				"w" ) as fp:
       fp.write(json.dumps(__A ) )
     with open(self.merges_file						,				"w" ) as fp:
       fp.write("\n".join(__A ) )






   def _snake_case    (		self						,				__A ):






     """simple docstring"""





     lowerCamelCase :					Dict    =				"lower newer"
     lowerCamelCase :					Union[str, Any]    =				"lower newer"
     return input_text, output_text






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					List[str]    =				BioGptTokenizer(self.vocab_file						,				self.merges_file )

     lowerCamelCase :					Optional[int]    =				"lower"
     lowerCamelCase :					Any    =				["low", "er</w>"]
     lowerCamelCase :					List[str]    =				tokenizer.tokenize(__A )
     self.assertListEqual(__A						,				__A )

     lowerCamelCase :					Union[str, Any]    =				tokens + ["<unk>"]
     lowerCamelCase :					List[str]    =				[14, 15, 20]
     self.assertListEqual(tokenizer.convert_tokens_to_ids(__A )						,				__A )






   @slow
   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					List[str]    =				BioGptTokenizer.from_pretrained("microsoft/biogpt" )

     lowerCamelCase :					Optional[int]    =				tokenizer.encode("sequence builders"						,				add_special_tokens=__A )
     lowerCamelCase :					Tuple    =				tokenizer.encode("multi-sequence build"						,				add_special_tokens=__A )

     lowerCamelCase :					Tuple    =				tokenizer.build_inputs_with_special_tokens(__A )
     lowerCamelCase :					List[str]    =				tokenizer.build_inputs_with_special_tokens(__A						,				__A )

     self.assertTrue(encoded_sentence == [2] + text )
     self.assertTrue(encoded_pair == [2] + text + [2] + text_a )



