




import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
		import torch

		from transformers.generation import (
		    MaxLengthCriteria,
		    MaxNewTokensCriteria,
		    MaxTimeCriteria,
		    StoppingCriteriaList,
		    validate_stopping_criteria,
		)


@require_torch
class 							__UpperCamelCase			( unittest.TestCase  ):





				'''simple docstring'''



				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		3
											lowerCAmelCase   =		2_50

											lowerCAmelCase   =		ids_tensor((batch_size, length)       , UpperCAmelCase_						)
											lowerCAmelCase   =		torch.ones((batch_size, length)       , device=UpperCAmelCase_       , dtype=torch.float						) / length
											return input_ids, scores



				def       __snake_case		( self						):
											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(5						)

											lowerCAmelCase   =		StoppingCriteriaList(
											    [
											        MaxLengthCriteria(max_length=10						),
											        MaxTimeCriteria(max_time=0.1						),
											    ]						)

											self.assertFalse(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(9						)
											self.assertFalse(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(10						)
											self.assertTrue(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		MaxLengthCriteria(max_length=10						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(5						)
											self.assertFalse(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(9						)
											self.assertFalse(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(10						)
											self.assertTrue(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		MaxNewTokensCriteria(start_length=5       , max_new_tokens=5						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(5						)
											self.assertFalse(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(9						)
											self.assertFalse(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(10						)
											self.assertTrue(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   =		StoppingCriteriaList([criteria]						)
											self.assertEqual(criteria_list.max_length       , 10						)



				def       __snake_case		( self						):
											lowerCAmelCase   ,     lowerCAmelCase   =		self._get_tensors(5						)

											lowerCAmelCase   =		MaxTimeCriteria(max_time=0.1						)
											self.assertFalse(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)

											lowerCAmelCase   =		MaxTimeCriteria(max_time=0.1       , initial_timestamp=time.time() - 0.2						)
											self.assertTrue(criteria(UpperCAmelCase_       , UpperCAmelCase_						)						)






				def       __snake_case		( self						):
											validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10						)]						)       , 10						)

											with self.assertWarns(UpperCAmelCase_						):
																		validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10						)]						)       , 11						)

											lowerCAmelCase   =		validate_stopping_criteria(StoppingCriteriaList()       , 11						)

											self.assertEqual(len(UpperCAmelCase_						)       , 1						)









import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
		import faiss


@require_faiss
class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				def       __snake_case		( self						):
											lowerCAmelCase   =		tempfile.mkdtemp()
											lowerCAmelCase   =		8

											# DPR tok
											lowerCAmelCase   =		[
											    '''[UNK]''',
											    '''[CLS]''',
											    '''[SEP]''',
											    '''[PAD]''',
											    '''[MASK]''',
											    '''want''',
											    '''##want''',
											    '''##ed''',
											    '''wa''',
											    '''un''',
											    '''runn''',
											    '''##ing''',
											    ''',''',
											    '''low''',
											    '''lowest''',
											]
											lowerCAmelCase   =		os.path.join(self.tmpdirname       , '''dpr_tokenizer'''						)
											os.makedirs(UpperCAmelCase_       , exist_ok=UpperCAmelCase_						)
											lowerCAmelCase   =		os.path.join(UpperCAmelCase_       , DPR_VOCAB_FILES_NAMES['''vocab_file''']						)
											with open(self.vocab_file       , '''w'''       , encoding='''utf-8'''						) as vocab_writer:
																		vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]						)						)

											# BART tok
											lowerCAmelCase   =		[
											    '''l''',
											    '''o''',
											    '''w''',
											    '''e''',
											    '''r''',
											    '''s''',
											    '''t''',
											    '''i''',
											    '''d''',
											    '''n''',
											    '''\u0120''',
											    '''\u0120l''',
											    '''\u0120n''',
											    '''\u0120lo''',
											    '''\u0120low''',
											    '''er''',
											    '''\u0120lowest''',
											    '''\u0120newer''',
											    '''\u0120wider''',
											    '''<unk>''',
											]
											lowerCAmelCase   =		dict(zip(UpperCAmelCase_       , range(len(UpperCAmelCase_						)						)						)						)
											lowerCAmelCase   =		['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
											lowerCAmelCase   =		{'''unk_token''': '''<unk>'''}

											lowerCAmelCase   =		os.path.join(self.tmpdirname       , '''bart_tokenizer'''						)
											os.makedirs(UpperCAmelCase_       , exist_ok=UpperCAmelCase_						)
											lowerCAmelCase   =		os.path.join(UpperCAmelCase_       , BART_VOCAB_FILES_NAMES['''vocab_file''']						)
											lowerCAmelCase   =		os.path.join(UpperCAmelCase_       , BART_VOCAB_FILES_NAMES['''merges_file''']						)
											with open(self.vocab_file       , '''w'''       , encoding='''utf-8'''						) as fp:
																		fp.write(json.dumps(UpperCAmelCase_						) + '''\n'''						)
											with open(self.merges_file       , '''w'''       , encoding='''utf-8'''						) as fp:
																		fp.write('''\n'''.join(UpperCAmelCase_						)						)



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



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



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



				def       __snake_case		( self						):
											shutil.rmtree(self.tmpdirname						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		Dataset.from_dict(
											    {
											        '''id''': ['''0''', '''1'''],
											        '''text''': ['''foo''', '''bar'''],
											        '''title''': ['''Foo''', '''Bar'''],
											        '''embeddings''': [np.ones(self.retrieval_vector_size						), 2 * np.ones(self.retrieval_vector_size						)],
											    }						)
											dataset.add_faiss_index('''embeddings'''       , string_factory='''Flat'''       , metric_type=faiss.METRIC_INNER_PRODUCT						)
											return dataset



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dummy_dataset()
											lowerCAmelCase   =		RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       , question_encoder=DPRConfig().to_dict()       , generator=BartConfig().to_dict()       , )
											with patch('''transformers.models.rag.retrieval_rag.load_dataset'''						) as mock_load_dataset:
																		lowerCAmelCase   =		dataset
																		lowerCAmelCase   =		RagRetriever(
																		    UpperCAmelCase_       , question_encoder_tokenizer=self.get_dpr_tokenizer()       , generator_tokenizer=self.get_bart_tokenizer()       , )
											return retriever



				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.get_dummy_dataset()
											lowerCAmelCase   =		RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size       , question_encoder=DPRConfig().to_dict()       , generator=BartConfig().to_dict()       , index_name='''custom'''       , )
											if from_disk:
																		lowerCAmelCase   =		os.path.join(self.tmpdirname       , '''dataset'''						)
																		lowerCAmelCase   =		os.path.join(self.tmpdirname       , '''index.faiss'''						)
																		dataset.get_index('''embeddings'''						).save(os.path.join(self.tmpdirname       , '''index.faiss'''						)						)
																		dataset.drop_index('''embeddings'''						)
																		dataset.save_to_disk(os.path.join(self.tmpdirname       , '''dataset'''						)						)
																		del dataset
																		lowerCAmelCase   =		RagRetriever(
																		    UpperCAmelCase_       , question_encoder_tokenizer=self.get_dpr_tokenizer()       , generator_tokenizer=self.get_bart_tokenizer()       , )
											else:
																		lowerCAmelCase   =		RagRetriever(
																		    UpperCAmelCase_       , question_encoder_tokenizer=self.get_dpr_tokenizer()       , generator_tokenizer=self.get_bart_tokenizer()       , index=CustomHFIndex(config.retrieval_vector_size       , UpperCAmelCase_						)       , )
											return retriever



				def       __snake_case		( self						):
											lowerCAmelCase   =		Dataset.from_dict(
											    {
											        '''id''': ['''0''', '''1'''],
											        '''text''': ['''foo''', '''bar'''],
											        '''title''': ['''Foo''', '''Bar'''],
											        '''embeddings''': [np.ones(self.retrieval_vector_size + 1						), 2 * np.ones(self.retrieval_vector_size + 1						)],
											    }						)
											dataset.add_faiss_index('''embeddings'''       , string_factory='''Flat'''       , metric_type=faiss.METRIC_INNER_PRODUCT						)

											lowerCAmelCase   =		os.path.join(self.tmpdirname       , '''hf_bert_base.hnswSQ8_correct_phi_128.c_index'''						)
											dataset.save_faiss_index('''embeddings'''       , index_file_name + '''.index.dpr'''						)
											pickle.dump(dataset['''id''']       , open(index_file_name + '''.index_meta.dpr'''       , '''wb'''						)						)

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

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



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_canonical_hf_index_retriever()
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''embeddings''', '''id''', '''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''id''']						)       , UpperCAmelCase_						)
											self.assertEqual(doc_dicts[0]['''id'''][0]       , '''1'''						)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['''id'''][0]       , '''0'''						)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist()       , [[1], [0]]						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dummy_canonical_hf_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		with patch('''transformers.models.rag.retrieval_rag.load_dataset'''						) as mock_load_dataset:
																									lowerCAmelCase   =		self.get_dummy_dataset()
																									retriever.save_pretrained(UpperCAmelCase_						)
																									lowerCAmelCase   =		RagRetriever.from_pretrained(UpperCAmelCase_						)
																									self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_						)
																									lowerCAmelCase   =		np.array(
																									    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
																									lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=1						)
																									self.assertTrue(out is not None						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''embeddings''', '''id''', '''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''id''']						)       , UpperCAmelCase_						)
											self.assertEqual(doc_dicts[0]['''id'''][0]       , '''1'''						)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['''id'''][0]       , '''0'''						)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist()       , [[1], [0]]						)



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



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''embeddings''', '''id''', '''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''id''']						)       , UpperCAmelCase_						)
											self.assertEqual(doc_dicts[0]['''id'''][0]       , '''1'''						)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['''id'''][0]       , '''0'''						)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist()       , [[1], [0]]						)



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



				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_legacy_index_retriever()
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		retriever.retrieve(UpperCAmelCase_       , n_docs=UpperCAmelCase_						)
											self.assertEqual(retrieved_doc_embeds.shape       , (2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(UpperCAmelCase_						)       , 2						)
											self.assertEqual(sorted(doc_dicts[0]						)       , ['''text''', '''title''']						)
											self.assertEqual(len(doc_dicts[0]['''text''']						)       , UpperCAmelCase_						)
											self.assertEqual(doc_dicts[0]['''text'''][0]       , '''bar'''						)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['''text'''][0]       , '''foo'''						)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist()       , [[1], [0]]						)



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



				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def       __snake_case		( self						):
											import torch

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

											lowerCAmelCase   =		retriever(
											    UpperCAmelCase_       , UpperCAmelCase_       , prefix=retriever.config.generator.prefix       , n_docs=UpperCAmelCase_       , return_tensors='''pt'''       , )
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		(  # 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(UpperCAmelCase_       , torch.Tensor						)
											self.assertIsInstance(UpperCAmelCase_       , torch.Tensor						)
											self.assertIsInstance(UpperCAmelCase_       , torch.Tensor						)






				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def       __snake_case		( self						):
											lowerCAmelCase   =		self.get_dpr_ctx_encoder_tokenizer()
											lowerCAmelCase   =		1
											lowerCAmelCase   =		self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_						)
											retriever.set_ctx_encoder_tokenizer(UpperCAmelCase_						)

											lowerCAmelCase   =		[[5, 7], [10, 11]]
											lowerCAmelCase   =		np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]       , dtype=np.floataa						)
											lowerCAmelCase   =		retriever(UpperCAmelCase_       , UpperCAmelCase_       , prefix=retriever.config.generator.prefix       , n_docs=UpperCAmelCase_						)

											self.assertEqual(
											    len(UpperCAmelCase_						)       , 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''')						)       , UpperCAmelCase_						)  # check for doc token related keys in dictionary.



