



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 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= tempfile.mkdtemp()
													__SCREAMING_SNAKE_CASE									= 8

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

													__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'bart_tokenizer'		)
													os.makedirs(_A				,							exist_ok=_A		)
													__SCREAMING_SNAKE_CASE									= os.path.join(_A				,							BART_VOCAB_FILES_NAMES['vocab_file']		)
													__SCREAMING_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 	_A     (   self		):







													'''simple docstring'''

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






						def 	_A     (   self		):







													'''simple docstring'''

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






						def 	_A     (   self		):







													'''simple docstring'''

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






						def 	_A     (   self		):







													'''simple docstring'''

													shutil.rmtree(self.tmpdirname		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_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 	_A     (   self		):







													'''simple docstring'''

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






						def 	_A     (   self				,							_A		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dummy_dataset()
													__SCREAMING_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:
																				__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'dataset'		)
																				__SCREAMING_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
																				__SCREAMING_SNAKE_CASE									= RagRetriever(
																				    _A				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer()				,							)
													else:
																				__SCREAMING_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 	_A     (   self		):







													'''simple docstring'''

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

													__SCREAMING_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'		)		)

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

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






						def 	_A     (   self		):







													'''simple docstring'''

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







													'''simple docstring'''

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






						def 	_A     (   self		):







													'''simple docstring'''

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







													'''simple docstring'''

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






						def 	_A     (   self		):







													'''simple docstring'''

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







													'''simple docstring'''

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






						def 	_A     (   self		):







													'''simple docstring'''

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







													'''simple docstring'''

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






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







													'''simple docstring'''

													import torch

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

													__SCREAMING_SNAKE_CASE									= retriever(
													    _A				,							_A				,							prefix=retriever.config.generator.prefix				,							n_docs=_A				,							return_tensors='pt'				,							)
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_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 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dpr_ctx_encoder_tokenizer()
													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= self.get_dummy_custom_hf_index_retriever(from_disk=_A		)
													retriever.set_ctx_encoder_tokenizer(_A		)

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

from transformers import LayoutLMTokenizer, LayoutLMTokenizerFast
from transformers.models.layoutlm.tokenization_layoutlm import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin

@require_tokenizers
class        UpperCAmelCase_  (       UpperCamelCase_   ,   unittest.TestCase	):

						'''simple docstring'''




						UpperCamelCase__      :			Optional[int]      =  LayoutLMTokenizer
						UpperCamelCase__      :			Any      =  LayoutLMTokenizerFast
						UpperCamelCase__      :			Optional[int]      =  True
						UpperCamelCase__      :			int      =  True






						def 	_A     (   self		):







													'''simple docstring'''

													super().setUp()

													__SCREAMING_SNAKE_CASE									= [
													    '[UNK]',
													    '[CLS]',
													    '[SEP]',
													    'want',
													    '##want',
													    '##ed',
													    'wa',
													    'un',
													    'runn',
													    '##ing',
													    ',',
													    'low',
													    'lowest',
													]
													__SCREAMING_SNAKE_CASE									= 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]		)		)






						def 	_A     (   self				,							**_A		):







													'''simple docstring'''

													return LayoutLMTokenizer.from_pretrained(self.tmpdirname				,							**_A		)






						def 	_A     (   self				,							_A		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= 'UNwant\u00E9d,running'
													__SCREAMING_SNAKE_CASE									= 'unwanted, running'
													return input_text, output_text






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.tokenizer_class(self.vocab_file		)

													__SCREAMING_SNAKE_CASE									= tokenizer.tokenize('UNwant\u00E9d,running'		)
													self.assertListEqual(_A				,							['un', '##want', '##ed', ',', 'runn', '##ing']		)
													self.assertListEqual(tokenizer.convert_tokens_to_ids(_A		)				,							[7, 4, 5, 10, 8, 9]		)






						def 	_A     (   self		):







													'''simple docstring'''

													pass






