
"""simple docstring"""


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

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
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_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
			from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
			__A 		= None

if is_ray_available():
			import ray  # noqa: E402 # isort:skip
			from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
			__A 		= None
			__A 		= None
			__A 		= None



def 					_lowerCamelCase(__UpperCamelCase )       ->				List[Any]:
	if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
		_lowerCAmelCase			=unittest.skip("""test requires Datasets, Faiss, psutil""" )(__UpperCamelCase )
	return test_case







@require_distributed_retrieval
class 	lowerCamelCase__      (      __magic_name__       ):





			'''simple docstring'''



			def       _lowerCAmelCase	(							self						)				->  int:
				_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       _lowerCAmelCase	(							self						)				->  DPRQuestionEncoderTokenizer:
				return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,  """dpr_tokenizer"""						)						)



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



			def       _lowerCAmelCase	(							self						)				->  int:
				shutil.rmtree(self.tmpdirname						)



			def       _lowerCAmelCase	(							self						)				->  str:
				_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       _lowerCAmelCase	(							self  ,  __UpperCAmelCase  ,  __UpperCAmelCase=1_23_45						)				->  RagPyTorchDistributedRetriever:
				_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			=RagPyTorchDistributedRetriever(
					    __UpperCAmelCase  ,  question_encoder_tokenizer=self.get_dpr_tokenizer()  ,  generator_tokenizer=self.get_bart_tokenizer()  ,  )
					if init_retrieval:
						retriever.init_retrieval(__UpperCAmelCase						)
				return retriever



			def       _lowerCAmelCase	(							self  ,  __UpperCAmelCase						)				->  RagRayDistributedRetriever:
				# Have to run in local mode because sys.path modifications at top of
				# file are not propogated to remote workers.
				# https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
				ray.init(local_mode=__UpperCAmelCase						)
				_lowerCAmelCase			=RagConfig(
				    retrieval_vector_size=self.retrieval_vector_size  ,  question_encoder=DPRConfig().to_dict()  ,  generator=BartConfig().to_dict()  ,  )
				_lowerCAmelCase			=ray.remote(__UpperCAmelCase						)
				_lowerCAmelCase			=[remote_cls.remote() for _ in range(1						)]
				with patch("""transformers.models.rag.retrieval_rag.load_dataset"""						) as mock_load_dataset:
					_lowerCAmelCase			=self.get_dummy_dataset()
					_lowerCAmelCase			=RagRayDistributedRetriever(
					    __UpperCAmelCase  ,  question_encoder_tokenizer=self.get_dpr_tokenizer()  ,  generator_tokenizer=self.get_bart_tokenizer()  ,  retrieval_workers=__UpperCAmelCase  ,  )
					if init_retrieval:
						retriever.init_retrieval()
				return retriever



			def       _lowerCAmelCase	(							self  ,  __UpperCAmelCase  ,  __UpperCAmelCase  ,  __UpperCAmelCase=1_23_45						)				->  Union[str, Any]:
				_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			=RagPyTorchDistributedRetriever(
					    __UpperCAmelCase  ,  question_encoder_tokenizer=self.get_dpr_tokenizer()  ,  generator_tokenizer=self.get_bart_tokenizer()  ,  )
				else:
					_lowerCAmelCase			=RagPyTorchDistributedRetriever(
					    __UpperCAmelCase  ,  question_encoder_tokenizer=self.get_dpr_tokenizer()  ,  generator_tokenizer=self.get_bart_tokenizer()  ,  index=CustomHFIndex(config.retrieval_vector_size  ,  __UpperCAmelCase						)  ,  )
				if init_retrieval:
					retriever.init_retrieval(__UpperCAmelCase						)
				return retriever



			def       _lowerCAmelCase	(							self  ,  __UpperCAmelCase  ,  __UpperCAmelCase						)				->  Union[str, Any]:
				# Have to run in local mode because sys.path modifications at top of
				# file are not propogated to remote workers.
				# https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
				ray.init(local_mode=__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"""  ,  )
				_lowerCAmelCase			=ray.remote(__UpperCAmelCase						)
				_lowerCAmelCase			=[remote_cls.remote() for _ in range(1						)]
				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			=RagRayDistributedRetriever(
					    __UpperCAmelCase  ,  question_encoder_tokenizer=self.get_dpr_tokenizer()  ,  generator_tokenizer=self.get_bart_tokenizer()  ,  retrieval_workers=__UpperCAmelCase  ,  index=CustomHFIndex.load_from_disk(
					        vector_size=config.retrieval_vector_size  ,  dataset_path=config.passages_path  ,  index_path=config.index_path  ,  )  ,  )
				else:
					_lowerCAmelCase			=RagRayDistributedRetriever(
					    __UpperCAmelCase  ,  question_encoder_tokenizer=self.get_dpr_tokenizer()  ,  generator_tokenizer=self.get_bart_tokenizer()  ,  retrieval_workers=__UpperCAmelCase  ,  index=CustomHFIndex(config.retrieval_vector_size  ,  __UpperCAmelCase						)  ,  )
				if init_retrieval:
					retriever.init_retrieval()
				return retriever



			def       _lowerCAmelCase	(							self  ,  __UpperCAmelCase  ,  __UpperCAmelCase  ,  __UpperCAmelCase						)				->  None:
				_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       _lowerCAmelCase	(							self						)				->  Any:
				_lowerCAmelCase			=1
				_lowerCAmelCase			=np.array(
				    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]  ,  dtype=np.floataa						)

				self.distributed_retriever_check(
				    self.get_dummy_pytorch_distributed_retriever(init_retrieval=__UpperCAmelCase						)  ,  __UpperCAmelCase  ,  __UpperCAmelCase						)



			def       _lowerCAmelCase	(							self						)				->  Union[str, Any]:
				_lowerCAmelCase			=1
				_lowerCAmelCase			=np.array(
				    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]  ,  dtype=np.floataa						)

				self.distributed_retriever_check(
				    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=__UpperCAmelCase  ,  from_disk=__UpperCAmelCase						)  ,  __UpperCAmelCase  ,  __UpperCAmelCase  ,  )



			def       _lowerCAmelCase	(							self						)				->  List[Any]:
				_lowerCAmelCase			=1
				_lowerCAmelCase			=np.array(
				    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]  ,  dtype=np.floataa						)

				self.distributed_retriever_check(
				    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=__UpperCAmelCase  ,  from_disk=__UpperCAmelCase						)  ,  __UpperCAmelCase  ,  __UpperCAmelCase  ,  )



			@require_ray
			def       _lowerCAmelCase	(							self						)				->  Tuple:
				_lowerCAmelCase			=1
				_lowerCAmelCase			=np.array(
				    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]  ,  dtype=np.floataa						)

				self.distributed_retriever_check(
				    self.get_dummy_ray_distributed_retriever(init_retrieval=__UpperCAmelCase						)  ,  __UpperCAmelCase  ,  __UpperCAmelCase						)
				ray.shutdown()



			@require_ray
			def       _lowerCAmelCase	(							self						)				->  Optional[Any]:
				_lowerCAmelCase			=1
				_lowerCAmelCase			=np.array(
				    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]  ,  dtype=np.floataa						)
				with self.assertRaises(__UpperCAmelCase						):
					self.distributed_retriever_check(
					    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=__UpperCAmelCase  ,  from_disk=__UpperCAmelCase						)  ,  __UpperCAmelCase  ,  __UpperCAmelCase  ,  )
				ray.shutdown()



			@require_ray
			def       _lowerCAmelCase	(							self						)				->  Optional[int]:
				_lowerCAmelCase			=1
				_lowerCAmelCase			=np.array(
				    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]  ,  dtype=np.floataa						)

				self.distributed_retriever_check(
				    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=__UpperCAmelCase  ,  from_disk=__UpperCAmelCase						)  ,  __UpperCAmelCase  ,  __UpperCAmelCase						)
				ray.shutdown()

"""simple docstring"""


from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A 		= logging.get_logger(__name__)

__A 		= {
    'nghuyong/ernie-1.0-base-zh': 'https://huggingface.co/nghuyong/ernie-1.0-base-zh/resolve/main/config.json',
    'nghuyong/ernie-2.0-base-en': 'https://huggingface.co/nghuyong/ernie-2.0-base-en/resolve/main/config.json',
    'nghuyong/ernie-2.0-large-en': 'https://huggingface.co/nghuyong/ernie-2.0-large-en/resolve/main/config.json',
    'nghuyong/ernie-3.0-base-zh': 'https://huggingface.co/nghuyong/ernie-3.0-base-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-medium-zh': 'https://huggingface.co/nghuyong/ernie-3.0-medium-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-mini-zh': 'https://huggingface.co/nghuyong/ernie-3.0-mini-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-micro-zh': 'https://huggingface.co/nghuyong/ernie-3.0-micro-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-nano-zh': 'https://huggingface.co/nghuyong/ernie-3.0-nano-zh/resolve/main/config.json',
    'nghuyong/ernie-gram-zh': 'https://huggingface.co/nghuyong/ernie-gram-zh/resolve/main/config.json',
    'nghuyong/ernie-health-zh': 'https://huggingface.co/nghuyong/ernie-health-zh/resolve/main/config.json',
}







class 	lowerCamelCase__      (      __magic_name__       ):





			'''simple docstring'''
			lowerCamelCase						      =	'''ernie'''



			def __init__(							self  ,  __UpperCAmelCase=3_05_22  ,  __UpperCAmelCase=7_68  ,  __UpperCAmelCase=12  ,  __UpperCAmelCase=12  ,  __UpperCAmelCase=30_72  ,  __UpperCAmelCase="gelu"  ,  __UpperCAmelCase=0.1  ,  __UpperCAmelCase=0.1  ,  __UpperCAmelCase=5_12  ,  __UpperCAmelCase=2  ,  __UpperCAmelCase=3  ,  __UpperCAmelCase=False  ,  __UpperCAmelCase=0.0_2  ,  __UpperCAmelCase=1e-12  ,  __UpperCAmelCase=0  ,  __UpperCAmelCase="absolute"  ,  __UpperCAmelCase=True  ,  __UpperCAmelCase=None  ,  **__UpperCAmelCase  ,  )				->  str:
				super().__init__(pad_token_id=__UpperCAmelCase  ,  **__UpperCAmelCase						)

				_lowerCAmelCase			=vocab_size
				_lowerCAmelCase			=hidden_size
				_lowerCAmelCase			=num_hidden_layers
				_lowerCAmelCase			=num_attention_heads
				_lowerCAmelCase			=hidden_act
				_lowerCAmelCase			=intermediate_size
				_lowerCAmelCase			=hidden_dropout_prob
				_lowerCAmelCase			=attention_probs_dropout_prob
				_lowerCAmelCase			=max_position_embeddings
				_lowerCAmelCase			=type_vocab_size
				_lowerCAmelCase			=task_type_vocab_size
				_lowerCAmelCase			=use_task_id
				_lowerCAmelCase			=initializer_range
				_lowerCAmelCase			=layer_norm_eps
				_lowerCAmelCase			=position_embedding_type
				_lowerCAmelCase			=use_cache
				_lowerCAmelCase			=classifier_dropout







class 	lowerCamelCase__      (      __magic_name__       ):





			'''simple docstring'''



			@property
			def       _lowerCAmelCase	(							self						)				->  Mapping[str, Mapping[int, str]]:
				if self.task == "multiple-choice":
					_lowerCAmelCase			={0: """batch""", 1: """choice""", 2: """sequence"""}
				else:
					_lowerCAmelCase			={0: """batch""", 1: """sequence"""}
				return OrderedDict(
				    [
				        ("""input_ids""", dynamic_axis),
				        ("""attention_mask""", dynamic_axis),
				        ("""token_type_ids""", dynamic_axis),
				        ("""task_type_ids""", dynamic_axis),
				    ]						)



