


'''simple docstring'''




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



	def 				lowercase     (					self    ):
		lowerCAmelCase					:    Optional[Any]      							= tempfile.mkdtemp()
		lowerCAmelCase					:    str      							= 8

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

		lowerCAmelCase					:    Tuple      							= os.path.join(self.tmpdirname ,					'bart_tokenizer'    )
		os.makedirs(UpperCAmelCase_ ,					exist_ok=UpperCAmelCase_    )
		lowerCAmelCase					:    Tuple      							= os.path.join(UpperCAmelCase_ ,					BART_VOCAB_FILES_NAMES['vocab_file']    )
		lowerCAmelCase					:    Tuple      							= 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 				lowercase     (					self    ):
		return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,					'dpr_tokenizer'    )    )



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



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



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



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



	def 				lowercase     (					self ,					snake_case__    ):
		lowerCAmelCase					:    List[str]      							= self.get_dummy_dataset()
		lowerCAmelCase					:    Tuple      							= RagConfig(
		    retrieval_vector_size=self.retrieval_vector_size ,					question_encoder=DPRConfig().to_dict() ,					generator=BartConfig().to_dict() ,					index_name='custom' ,					)
		if from_disk:
			lowerCAmelCase					:    int      							= os.path.join(self.tmpdirname ,					'dataset'    )
			lowerCAmelCase					:    int      							= 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					:    Any      							= RagRetriever(
			    UpperCAmelCase_ ,					question_encoder_tokenizer=self.get_dpr_tokenizer() ,					generator_tokenizer=self.get_bart_tokenizer() ,					)
		else:
			lowerCAmelCase					:    Tuple      							= 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 				lowercase     (					self    ):
		lowerCAmelCase					:    Union[str, Any]      							= 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					:    Union[str, Any]      							= 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					:    Dict      							= os.path.join(self.tmpdirname ,					'psgs_w100.tsv.pkl'    )
		lowerCAmelCase					:    int      							= {sample['id']: [sample['text'], sample['title']] for sample in dataset}
		pickle.dump(UpperCAmelCase_ ,					open(UpperCAmelCase_ ,					'wb'    )    )

		lowerCAmelCase					:    Optional[Any]      							= 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					:    str      							= RagRetriever(
		    UpperCAmelCase_ ,					question_encoder_tokenizer=self.get_dpr_tokenizer() ,					generator_tokenizer=self.get_bart_tokenizer()    )
		return retriever



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



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



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



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



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

		lowerCAmelCase					:    Optional[int]      							= 1
		lowerCAmelCase					:    Union[str, Any]      							= self.get_dummy_canonical_hf_index_retriever()
		lowerCAmelCase					:    Any      							= [[5, 7], [10, 11]]
		lowerCAmelCase					:    Optional[Any]      							= np.array(
		    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )] ,					dtype=np.floataa    )
		lowerCAmelCase					:    List[Any]      							= retriever(UpperCAmelCase_ ,					UpperCAmelCase_ ,					prefix=retriever.config.generator.prefix ,					n_docs=UpperCAmelCase_    )
		lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					:    Optional[Any]      							= (
		    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					:    Tuple      							= retriever(
		    UpperCAmelCase_ ,					UpperCAmelCase_ ,					prefix=retriever.config.generator.prefix ,					n_docs=UpperCAmelCase_ ,					return_tensors='pt' ,					)
		lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					:    int      							= (  # 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 				lowercase     (					self    ):
		lowerCAmelCase					:    int      							= self.get_dpr_ctx_encoder_tokenizer()
		lowerCAmelCase					:    List[Any]      							= 1
		lowerCAmelCase					:    str      							= self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_    )
		retriever.set_ctx_encoder_tokenizer(UpperCAmelCase_    )

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



'''simple docstring'''




import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict

@pytest.mark.parametrize(
    'files'					,   [
        ['full:README.md', 'dataset_infos.json'],
        ['empty:README.md', 'dataset_infos.json'],
        ['dataset_infos.json'],
        ['full:README.md'],
    ]					,   )
def 			__UpperCamelCase    (    _A    :				str					,   _A    :				List[Any]				)    ->      Union[str, Any]:




	"""simple docstring"""





	lowerCAmelCase					:    Optional[int]      							= tmp_path_factory.mktemp('dset_infos_dir'				)
	if "full:README.md" in files:
		with open(dataset_infos_dir / 'README.md'					,   'w'				) as f:
			f.write('---\ndataset_info:\n  dataset_size: 42\n---'				)
	if "empty:README.md" in files:
		with open(dataset_infos_dir / 'README.md'					,   'w'				) as f:
			f.write(''				)
    # we want to support dataset_infos.json for backward compatibility
	if "dataset_infos.json" in files:
		with open(dataset_infos_dir / 'dataset_infos.json'					,   'w'				) as f:
			f.write('{"default": {"dataset_size": 42}}'				)
	lowerCAmelCase					:    Union[str, Any]      							= DatasetInfosDict.from_directory(_A				)
	assert dataset_infos
	assert dataset_infos["default"].dataset_size == 42

@pytest.mark.parametrize(
    'dataset_info'					,   [
        DatasetInfo(),
        DatasetInfo(
            description='foo'					,   features=Features({'a': Value('int32'				)}				)					,   builder_name='builder'					,   config_name='config'					,   version='1.0.0'					,   splits=[{'name': 'train'}]					,   download_size=42					,   ),
    ]					,   )
def 			__UpperCamelCase    (    _A    :				str					,   _A    :				DatasetInfo				)    ->      Optional[int]:




	"""simple docstring"""





	lowerCAmelCase					:    str      							= str(_A				)
	dataset_info.write_to_directory(_A				)
	lowerCAmelCase					:    List[str]      							= DatasetInfo.from_directory(_A				)
	assert dataset_info == reloaded
	assert os.path.exists(os.path.join(_A					,   'dataset_info.json'				)				)

def 			__UpperCamelCase    (    )    ->      List[str]:




	"""simple docstring"""





	lowerCAmelCase					:    Tuple      							= DatasetInfo(
	    description='foo'					,   citation='bar'					,   homepage='https://foo.bar'					,   license='CC0'					,   features=Features({'a': Value('int32'				)}				)					,   post_processed={}					,   supervised_keys=()					,   task_templates=[]					,   builder_name='builder'					,   config_name='config'					,   version='1.0.0'					,   splits=[{'name': 'train', 'num_examples': 42}]					,   download_checksums={}					,   download_size=13_37					,   post_processing_size=4_42					,   dataset_size=12_34					,   size_in_bytes=13_37 + 4_42 + 12_34					,   )
	lowerCAmelCase					:    Optional[int]      							= dataset_info._to_yaml_dict()
	assert sorted(_A				) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML				)
	for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
		assert key in dataset_info_yaml_dict
		assert isinstance(dataset_info_yaml_dict[key]					,   (list, dict, int, str)				)
	lowerCAmelCase					:    Any      							= yaml.safe_dump(_A				)
	lowerCAmelCase					:    int      							= yaml.safe_load(_A				)
	assert dataset_info_yaml_dict == reloaded

def 			__UpperCamelCase    (    )    ->      Dict:




	"""simple docstring"""





	lowerCAmelCase					:    Union[str, Any]      							= DatasetInfo()
	lowerCAmelCase					:    List[Any]      							= dataset_info._to_yaml_dict()
	assert dataset_info_yaml_dict == {}






@pytest.mark.parametrize(
    'dataset_infos_dict'					,   [
        DatasetInfosDict(),
        DatasetInfosDict({'default': DatasetInfo()}				),
        DatasetInfosDict({'my_config_name': DatasetInfo()}				),
        DatasetInfosDict(
            {
                'default': DatasetInfo(
                    description='foo'					,   features=Features({'a': Value('int32'				)}				)					,   builder_name='builder'					,   config_name='config'					,   version='1.0.0'					,   splits=[{'name': 'train'}]					,   download_size=42					,   )
            }				),
        DatasetInfosDict(
            {
                'v1': DatasetInfo(dataset_size=42				),
                'v2': DatasetInfo(dataset_size=13_37				),
            }				),
    ]					,   )
def 			__UpperCamelCase    (    _A    :				Tuple					,   _A    :				DatasetInfosDict				)    ->      List[Any]:




	"""simple docstring"""





	lowerCAmelCase					:    Tuple      							= str(_A				)
	dataset_infos_dict.write_to_directory(_A				)
	lowerCAmelCase					:    List[str]      							= DatasetInfosDict.from_directory(_A				)

	# the config_name of the dataset_infos_dict take over the attribute
	for config_name, dataset_info in dataset_infos_dict.items():
		lowerCAmelCase					:    Tuple      							= config_name
		# the yaml representation doesn't include fields like description or citation
		# so we just test that we can recover what we can from the yaml
		lowerCAmelCase					:    Optional[Any]      							= DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()				)
	assert dataset_infos_dict == reloaded

	if dataset_infos_dict:
		assert os.path.exists(os.path.join(_A					,   'README.md'				)				)
