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





     def 		UpperCAmelCase(self				:   int       )						->						Tuple:
         snake_case								= tempfile.mkdtemp()
         snake_case								= 8

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

         snake_case								= os.path.join(self.tmpdirname    ,   "bart_tokenizer"       )
         os.makedirs(_A    ,   exist_ok=_A       )
         snake_case								= os.path.join(_A    ,   BART_VOCAB_FILES_NAMES["vocab_file"]       )
         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 		UpperCAmelCase(self				:   Optional[int]       )						->						DPRQuestionEncoderTokenizer:
         return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname    ,   "dpr_tokenizer"       )       )





     def 		UpperCAmelCase(self				:   List[Any]       )						->						DPRContextEncoderTokenizer:
         return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname    ,   "dpr_tokenizer"       )       )





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





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





     def 		UpperCAmelCase(self				:   Optional[Any]       )						->						Dict:
         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 		UpperCAmelCase(self				:   Union[str, Any]       )						->						str:
         snake_case								= self.get_dummy_dataset()
         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:
             snake_case								= dataset
             snake_case								= RagRetriever(
                 _A    ,   question_encoder_tokenizer=self.get_dpr_tokenizer()    ,   generator_tokenizer=self.get_bart_tokenizer()    ,   )
         return retriever





     def 		UpperCAmelCase(self				:   Dict    ,   _A				:   bool       )						->						List[str]:
         snake_case								= self.get_dummy_dataset()
         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:
             snake_case								= os.path.join(self.tmpdirname    ,   "dataset"       )
             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
             snake_case								= RagRetriever(
                 _A    ,   question_encoder_tokenizer=self.get_dpr_tokenizer()    ,   generator_tokenizer=self.get_bart_tokenizer()    ,   )
         else:
             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 		UpperCAmelCase(self				:   Optional[Any]       )						->						str:
         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       )

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

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

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





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





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





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





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





     @require_torch
     @require_tokenizers
     @require_sentencepiece
     def 		UpperCAmelCase(self				:   Dict       )						->						Dict:
         import torch

         snake_case								= 1
         snake_case								= self.get_dummy_canonical_hf_index_retriever()
         snake_case								= [[5, 7], [1_0, 1_1]]
         snake_case								= np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]    ,   dtype=np.floataa       )
         snake_case								= retriever(_A    ,   _A    ,   prefix=retriever.config.generator.prefix    ,   n_docs=_A       )
         snake_case					,							snake_case					,							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       )

         snake_case								= retriever(
             _A    ,   _A    ,   prefix=retriever.config.generator.prefix    ,   n_docs=_A    ,   return_tensors="pt"    ,   )
         snake_case					,							snake_case					,							snake_case					,							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 		UpperCAmelCase(self				:   Any       )						->						Any:
         snake_case								= self.get_dpr_ctx_encoder_tokenizer()
         snake_case								= 1
         snake_case								= self.get_dummy_custom_hf_index_retriever(from_disk=_A       )
         retriever.set_ctx_encoder_tokenizer(_A       )

         snake_case								= [[5, 7], [1_0, 1_1]]
         snake_case								= np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]    ,   dtype=np.floataa       )
         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.



import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_mobilebert import MobileBertTokenizer


_A   =			logging.get_logger(__name__)

_A   =			{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

_A   =			{
    "vocab_file": {"mobilebert-uncased": "https://huggingface.co/google/mobilebert-uncased/resolve/main/vocab.txt"},
    "tokenizer_file": {
        "mobilebert-uncased": "https://huggingface.co/google/mobilebert-uncased/resolve/main/tokenizer.json"
    },
}

_A   =			{"mobilebert-uncased": 5_12}


_A   =			{}







class 					lowerCamelCase		(  A_  ):
     UpperCAmelCase__    :							Tuple									=		VOCAB_FILES_NAMES
     UpperCAmelCase__    :							Any									=		PRETRAINED_VOCAB_FILES_MAP
     UpperCAmelCase__    :							Tuple									=		PRETRAINED_INIT_CONFIGURATION
     UpperCAmelCase__    :							Tuple									=		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     UpperCAmelCase__    :							Optional[Any]									=		MobileBertTokenizer





     def __init__(self				:   Any    ,   _A				:   str=None    ,   _A				:   str=None    ,   _A				:   Union[str, Any]=True    ,   _A				:   Optional[Any]="[UNK]"    ,   _A				:   int="[SEP]"    ,   _A				:   Dict="[PAD]"    ,   _A				:   int="[CLS]"    ,   _A				:   Union[str, Any]="[MASK]"    ,   _A				:   Any=True    ,   _A				:   Dict=None    ,   **_A				:   List[str]    ,   )						->						List[str]:
         super().__init__(
             _A    ,   tokenizer_file=_A    ,   do_lower_case=_A    ,   unk_token=_A    ,   sep_token=_A    ,   pad_token=_A    ,   cls_token=_A    ,   mask_token=_A    ,   tokenize_chinese_chars=_A    ,   strip_accents=_A    ,   **_A    ,   )

         snake_case								= json.loads(self.backend_tokenizer.normalizer.__getstate__()       )
         if (
             normalizer_state.get("lowercase"    ,   _A       ) != do_lower_case
             or normalizer_state.get("strip_accents"    ,   _A       ) != strip_accents
             or normalizer_state.get("handle_chinese_chars"    ,   _A       ) != tokenize_chinese_chars
         ):
             snake_case								= getattr(_A    ,   normalizer_state.pop("type"       )       )
             snake_case								= do_lower_case
             snake_case								= strip_accents
             snake_case								= tokenize_chinese_chars
             snake_case								= normalizer_class(**_A       )

         snake_case								= do_lower_case





     def 		UpperCAmelCase(self				:   List[str]    ,   _A				:   Union[str, Any]    ,   _A				:   Dict=None       )						->						Optional[Any]:
         snake_case								= [self.cls_token_id] + token_ids_a + [self.sep_token_id]

         if token_ids_a:
             output += token_ids_a + [self.sep_token_id]

         return output





     def 		UpperCAmelCase(self				:   Union[str, Any]    ,   _A				:   List[int]    ,   _A				:   Optional[List[int]] = None       )						->						List[int]:
         snake_case								= [self.sep_token_id]
         snake_case								= [self.cls_token_id]
         if token_ids_a is None:
             return len(cls + token_ids_a + sep       ) * [0]
         return len(cls + token_ids_a + sep       ) * [0] + len(token_ids_a + sep       ) * [1]





     def 		UpperCAmelCase(self				:   int    ,   _A				:   str    ,   _A				:   Optional[str] = None       )						->						Tuple[str]:
         snake_case								= self._tokenizer.model.save(_A    ,   name=_A       )
         return tuple(_A       )

