



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__(				lowerCamelCase						):

   '''simple docstring'''




   def       UpperCAmelCase		(    self	:      List[Any])						->					Optional[Any]:




         """simple docstring"""

         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(lowerCAmelCase      ,					exist_ok=lowerCAmelCase)
         lowercase__  =	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
         lowercase__  =	[
             'l',
             'o',
             'w',
             'e',
             'r',
             's',
             't',
             'i',
             'd',
             'n',
             '\u0120',
             '\u0120l',
             '\u0120n',
             '\u0120lo',
             '\u0120low',
             'er',
             '\u0120lowest',
             '\u0120newer',
             '\u0120wider',
             '<unk>',
         ]
         lowercase__  =	dict(zip(lowerCAmelCase      ,					range(len(lowerCAmelCase))))
         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(lowerCAmelCase      ,					exist_ok=lowerCAmelCase)
         lowercase__  =	os.path.join(lowerCAmelCase      ,					BART_VOCAB_FILES_NAMES['vocab_file'])
         lowercase__  =	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       UpperCAmelCase		(    self	:      Dict)						->					DPRQuestionEncoderTokenizer:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      int)						->					DPRContextEncoderTokenizer:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Any)						->					BartTokenizer:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Optional[Any])						->					List[str]:




         """simple docstring"""

         shutil.rmtree(self.tmpdirname)


   def       UpperCAmelCase		(    self	:      List[Any])						->					Optional[Any]:




         """simple docstring"""

         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       UpperCAmelCase		(    self	:      List[str])						->					Tuple:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Any      ,					lowerCAmelCase	:      bool)						->					List[Any]:




         """simple docstring"""

         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(
                   lowerCAmelCase      ,					question_encoder_tokenizer=self.get_dpr_tokenizer()      ,					generator_tokenizer=self.get_bart_tokenizer()      ,					)
         else:
               lowercase__  =	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       UpperCAmelCase		(    self	:      Any)						->					Any:




         """simple docstring"""

         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(lowerCAmelCase      ,					open(lowerCAmelCase      ,					'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(
             lowerCAmelCase      ,					question_encoder_tokenizer=self.get_dpr_tokenizer()      ,					generator_tokenizer=self.get_bart_tokenizer())
         return retriever


   def       UpperCAmelCase		(    self	:      str)						->					Dict:




         """simple docstring"""

         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(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       UpperCAmelCase		(    self	:      Any)						->					Dict:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      str)						->					Union[str, Any]:




         """simple docstring"""

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




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Optional[int])						->					Union[str, Any]:




         """simple docstring"""

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




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      List[Any])						->					List[Any]:




         """simple docstring"""

         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(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       UpperCAmelCase		(    self	:      Any)						->					int:




         """simple docstring"""

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


   @require_torch
   @require_tokenizers
   @require_sentencepiece
   def       UpperCAmelCase		(    self	:      int)						->					str:




         """simple docstring"""

         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(lowerCAmelCase      ,					lowerCAmelCase      ,					prefix=retriever.config.generator.prefix      ,					n_docs=lowerCAmelCase)
         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(lowerCAmelCase      ,					lowerCAmelCase)
         self.assertIsInstance(lowerCAmelCase      ,					lowerCAmelCase)
         self.assertIsInstance(lowerCAmelCase      ,					np.ndarray)

         lowercase__  =	retriever(
             lowerCAmelCase      ,					lowerCAmelCase      ,					prefix=retriever.config.generator.prefix      ,					n_docs=lowerCAmelCase      ,					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(lowerCAmelCase      ,					torch.Tensor)
         self.assertIsInstance(lowerCAmelCase      ,					torch.Tensor)
         self.assertIsInstance(lowerCAmelCase      ,					torch.Tensor)



   @require_torch
   @require_tokenizers
   @require_sentencepiece
   def       UpperCAmelCase		(    self	:      int)						->					str:




         """simple docstring"""

         lowercase__  =	self.get_dpr_ctx_encoder_tokenizer()
         lowercase__  =	1
         lowercase__  =	self.get_dummy_custom_hf_index_retriever(from_disk=lowerCAmelCase)
         retriever.set_ctx_encoder_tokenizer(lowerCAmelCase)

         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(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.





import os
import unittest

from transformers.models.cpmant.tokenization_cpmant import VOCAB_FILES_NAMES, CpmAntTokenizer
from transformers.testing_utils import require_jieba, tooslow

from ...test_tokenization_common import TokenizerTesterMixin



@require_jieba
class    UpperCAmelCase__(				lowerCamelCase  ,				unittest.TestCase						):

   '''simple docstring'''




   A     :	Optional[int]												=     CpmAntTokenizer
   A     :	Optional[int]												=     False


   def       UpperCAmelCase		(    self	:      List[str])						->					str:




         """simple docstring"""

         super().setUp()

         lowercase__  =	[
             '<d>',
             '</d>',
             '<s>',
             '</s>',
             '</_>',
             '<unk>',
             '<pad>',
             '</n>',
             '我',
             '是',
             'C',
             'P',
             'M',
             'A',
             'n',
             't',
         ]
         lowercase__  =	os.path.join(self.tmpdirname      ,					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]))



   @tooslow
   def       UpperCAmelCase		(    self	:      str)						->					Optional[Any]:




         """simple docstring"""

         lowercase__  =	CpmAntTokenizer.from_pretrained('openbmb/cpm-ant-10b')
         lowercase__  =	'今天天气真好！'
         lowercase__  =	['今天', '天气', '真', '好', '！']
         lowercase__  =	tokenizer.tokenize(lowerCAmelCase)
         self.assertListEqual(lowerCAmelCase      ,					lowerCAmelCase)
         lowercase__  =	'今天天气真好！'
         lowercase__  =	[tokenizer.bos_token] + tokens

         lowercase__  =	[6, 98_02, 1_49_62, 20_82, 8_31, 2_44]
         self.assertListEqual(tokenizer.convert_tokens_to_ids(lowerCAmelCase)      ,					lowerCAmelCase)

         lowercase__  =	tokenizer.decode(lowerCAmelCase)
         self.assertEqual(lowerCAmelCase      ,					lowerCAmelCase)
