



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



       def 							_a      (				self      )  ->							Any:
           lowercase__		       =						tempfile.mkdtemp()
           lowercase__		       =						8

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

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



       def 							_a      (				self      )  ->							DPRContextEncoderTokenizer:
           return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,"dpr_tokenizer"      )      )



       def 							_a      (				self      )  ->							BartTokenizer:
           return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,"bart_tokenizer"      )      )



       def 							_a      (				self      )  ->							List[Any]:
           shutil.rmtree(self.tmpdirname      )



       def 							_a      (				self      )  ->							Optional[Any]:
           lowercase__		       =						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      )  ->							str:
           lowercase__		       =						self.get_dummy_dataset()
           lowercase__		       =						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:
               lowercase__		       =						dataset
               lowercase__		       =						RagRetriever(
                   __A     ,question_encoder_tokenizer=self.get_dpr_tokenizer()     ,generator_tokenizer=self.get_bart_tokenizer()     ,)
           return retriever



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

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

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

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



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



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



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



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



       @require_torch
       @require_tokenizers
       @require_sentencepiece
       def 							_a      (				self      )  ->							Tuple:
           import torch

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

           lowercase__		       =						retriever(
               __A     ,__A     ,prefix=retriever.config.generator.prefix     ,n_docs=__A     ,return_tensors="pt"     ,)
           lowercase__       ,       lowercase__       ,       lowercase__       ,       lowercase__		       =						(  # 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 							_a      (				self      )  ->							int:
           lowercase__		       =						self.get_dpr_ctx_encoder_tokenizer()
           lowercase__		       =						1
           lowercase__		       =						self.get_dummy_custom_hf_index_retriever(from_disk=__A      )
           retriever.set_ctx_encoder_tokenizer(__A      )

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









'''simple docstring'''





import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger


class 				snake_case (UpperCamelCase   ):



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							int:
           lowercase__		       =						params

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array([len(UpperCAmelCase_      ) for t in data]      )

           self.check()
           self.remove_long_sequences()
           self.remove_empty_sequences()
           self.remove_unknown_sequences()
           self.check()
           self.print_statistics()



       def __getitem__(				self     ,UpperCAmelCase_      )  ->							Optional[int]:
           return (self.token_ids[index], self.lengths[index])



       def __len__(				self      )  ->							Union[str, Any]:
           return len(self.lengths      )



       def 							_a      (				self      )  ->							int:
           assert len(self.token_ids      ) == len(self.lengths      )
           assert all(self.lengths[i] == len(self.token_ids[i]      ) for i in range(len(self.lengths      )      )      )



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						self.params.max_model_input_size
           lowercase__		       =						self.lengths > max_len
           logger.info(F'''Splitting {sum(UpperCAmelCase_      )} too long sequences.'''      )

           def divide_chunks(UpperCAmelCase_     ,UpperCAmelCase_      ):
               return [l[i : i + n] for i in range(0     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )]

           lowercase__		       =						[]
           lowercase__		       =						[]
           if self.params.mlm:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["cls_token"], self.params.special_tok_ids["sep_token"]
           else:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["bos_token"], self.params.special_tok_ids["eos_token"]

           for seq_, len_ in zip(self.token_ids     ,self.lengths      ):
               assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
               if len_ <= max_len:
                   new_tok_ids.append(seq_      )
                   new_lengths.append(len_      )
               else:
                   lowercase__		       =						[]
                   for sub_s in divide_chunks(seq_     ,max_len - 2      ):
                       if sub_s[0] != cls_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,0     ,UpperCAmelCase_      )
                       if sub_s[-1] != sep_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )
                       assert len(UpperCAmelCase_      ) <= max_len
                       assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
                       sub_seqs.append(UpperCAmelCase_      )

                   new_tok_ids.extend(UpperCAmelCase_      )
                   new_lengths.extend([len(UpperCAmelCase_      ) for l in sub_seqs]      )

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array(UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Any:
           lowercase__		       =						len(self      )
           lowercase__		       =						self.lengths > 11
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} too short (<=11 tokens) sequences.'''      )



       def 							_a      (				self      )  ->							List[Any]:
           if "unk_token" not in self.params.special_tok_ids:
               return
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						len(self      )
           lowercase__		       =						np.array([np.count_nonzero(a == unk_token_id      ) for a in self.token_ids]      )
           lowercase__		       =						(unk_occs / self.lengths) < 0.5
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).'''      )



       def 							_a      (				self      )  ->							Optional[int]:
           if not self.params.is_master:
               return
           logger.info(F'''{len(self      )} sequences'''      )
           # data_len = sum(self.lengths)
           # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
           # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

           # unk_idx = self.params.special_tok_ids['unk_token']
           # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
           # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')







       def 							_a      (				self     ,UpperCAmelCase_      )  ->							List[str]:
           lowercase__		       =						[t[0] for t in batch]
           lowercase__		       =						[t[1] for t in batch]
           assert len(UpperCAmelCase_      ) == len(UpperCAmelCase_      )

           # Max for paddings
           lowercase__		       =						max(UpperCAmelCase_      )

           # Pad token ids
           if self.params.mlm:
               lowercase__		       =						self.params.special_tok_ids["pad_token"]
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						[list(t.astype(UpperCAmelCase_      )      ) + [pad_idx] * (max_seq_len_ - len(UpperCAmelCase_      )) for t in token_ids]
           assert len(tk_      ) == len(UpperCAmelCase_      )
           assert all(len(UpperCAmelCase_      ) == max_seq_len_ for t in tk_      )

           lowercase__		       =						torch.tensor(tk_      )  # (bs, max_seq_len_)
           lowercase__		       =						torch.tensor(UpperCAmelCase_      )  # (bs)
           return tk_t, lg_t




