


'''simple docstring'''






import json
import os
import unittest

from transformers.models.gpt_neox_japanese.tokenization_gpt_neox_japanese import (
    VOCAB_FILES_NAMES,
    GPTNeoXJapaneseTokenizer,
)
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin




@require_tokenizers
class 		a__				(							a__				,				unittest.TestCase  ):





							'''simple docstring'''






							lowercase__			:   Optional[Any]							=      GPTNeoXJapaneseTokenizer
							lowercase__			:   List[Any]							=      False
							lowercase__			:   Optional[int]							=      {"do_clean_text": False, "add_prefix_space": False}







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Any:
														super().setUp()

														lowerCAmelCase__          =	[
														    '''こん''',
														    '''こんに''',
														    '''にちは''',
														    '''ばんは''',
														    '''世界,㔺界''',
														    '''、''',
														    '''。''',
														    '''<BR>''',
														    '''<SP>''',
														    '''<TAB>''',
														    '''<URL>''',
														    '''<EMAIL>''',
														    '''<TEL>''',
														    '''<DATE>''',
														    '''<PRICE>''',
														    '''<BLOCK>''',
														    '''<KIGOU>''',
														    '''<U2000U2BFF>''',
														    '''<|emoji1|>''',
														    '''<unk>''',
														    '''<|startoftext|>''',
														    '''<|endoftext|>''',
														]
														lowerCAmelCase__          =	{'''emoji''': {'''\ud83d\ude00''': '''<|emoji1|>'''}, '''emoji_inv''': {'''<|emoji1|>''': '''\ud83d\ude00'''}}  # 😀
														lowerCAmelCase__          =	{'''unk_token''': '''<unk>'''}

														lowerCAmelCase__          =	os.path.join(self.tmpdirname					,    VOCAB_FILES_NAMES['''vocab_file''']   )
														lowerCAmelCase__          =	os.path.join(self.tmpdirname					,    VOCAB_FILES_NAMES['''emoji_file''']   )
														with open(self.vocab_file					,    '''w'''					,    encoding='''utf-8'''   ) as vocab_writer:
																					vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]   )   )
														with open(self.emoji_file					,    '''w'''   ) as emoji_writer:
																					emoji_writer.write(json.dumps(lowerCamelCase_   )   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    **lowerCamelCase_   )       ->							Optional[Any]:
														kwargs.update(self.special_tokens_map   )
														return GPTNeoXJapaneseTokenizer.from_pretrained(self.tmpdirname					,    **lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							Tuple:
														lowerCAmelCase__          =	'''こんにちは、世界。 \nこんばんは、㔺界。😀'''
														lowerCAmelCase__          =	'''こんにちは、世界。 \nこんばんは、世界。😀'''
														return input_text, output_text







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							Tuple:
														lowerCAmelCase__   ,					lowerCAmelCase__          =	self.get_input_output_texts(lowerCamelCase_   )
														lowerCAmelCase__          =	tokenizer.encode(lowerCamelCase_					,    add_special_tokens=lowerCamelCase_   )
														lowerCAmelCase__          =	tokenizer.decode(lowerCamelCase_					,    clean_up_tokenization_spaces=lowerCamelCase_   )
														return text, ids







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														pass  # TODO add if relevant







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[Any]:
														pass  # TODO add if relevant







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														pass  # TODO add if relevant







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[Any]:
														lowerCAmelCase__          =	self.get_tokenizer()

														# Testing tokenization
														lowerCAmelCase__          =	'''こんにちは、世界。　こんばんは、㔺界。'''
														lowerCAmelCase__          =	['''こん''', '''にちは''', '''、''', '''世界''', '''。''', '''<SP>''', '''こん''', '''ばんは''', '''、''', '''㔺界''', '''。''']
														lowerCAmelCase__          =	tokenizer.tokenize(lowerCamelCase_   )
														self.assertListEqual(lowerCamelCase_					,    lowerCamelCase_   )

														# Testing conversion to ids without special tokens
														lowerCAmelCase__          =	[0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6]
														lowerCAmelCase__          =	tokenizer.convert_tokens_to_ids(lowerCamelCase_   )
														self.assertListEqual(lowerCamelCase_					,    lowerCamelCase_   )

														# Testing conversion to ids with special tokens
														lowerCAmelCase__          =	tokens + [tokenizer.unk_token]
														lowerCAmelCase__          =	[0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6, 19]
														lowerCAmelCase__          =	tokenizer.convert_tokens_to_ids(lowerCamelCase_   )
														self.assertListEqual(lowerCamelCase_					,    lowerCamelCase_   )







							@slow
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[int]:
														lowerCAmelCase__          =	self.tokenizer_class.from_pretrained('''abeja/gpt-neox-japanese-2.7b'''   )

														lowerCAmelCase__          =	tokenizer.encode('''ありがとう。'''					,    add_special_tokens=lowerCamelCase_   )
														lowerCAmelCase__          =	tokenizer.encode('''どういたしまして。'''					,    add_special_tokens=lowerCamelCase_   )

														lowerCAmelCase__          =	tokenizer.build_inputs_with_special_tokens(lowerCamelCase_   )
														lowerCAmelCase__          =	tokenizer.build_inputs_with_special_tokens(lowerCamelCase_					,    lowerCamelCase_   )

														assert encoded_sentence == ids_a
														assert encoded_pair == ids_a + ids_a







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[str]:
														# Intentionally convert some words to accommodate character fluctuations unique to Japanese
														pass







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							str:
														# tokenizer has no padding token
														pass



'''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:
		__UpperCAmelCase    				=							None

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


def 					_snake_case   (     A       )						->    Any:
							if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
														lowerCAmelCase__          =	unittest.skip('''test requires Datasets, Faiss, psutil'''       )(A       )
							return test_case






@require_distributed_retrieval
class 		a__				(							a__  ):





							'''simple docstring'''






							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Union[str, Any]:
														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(lowerCamelCase_					,    exist_ok=lowerCamelCase_   )
														lowerCAmelCase__          =	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
														lowerCAmelCase__          =	[
														    '''l''',
														    '''o''',
														    '''w''',
														    '''e''',
														    '''r''',
														    '''s''',
														    '''t''',
														    '''i''',
														    '''d''',
														    '''n''',
														    '''\u0120''',
														    '''\u0120l''',
														    '''\u0120n''',
														    '''\u0120lo''',
														    '''\u0120low''',
														    '''er''',
														    '''\u0120lowest''',
														    '''\u0120newer''',
														    '''\u0120wider''',
														    '''<unk>''',
														]
														lowerCAmelCase__          =	dict(zip(lowerCamelCase_					,    range(len(lowerCamelCase_   )   )   )   )
														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(lowerCamelCase_					,    exist_ok=lowerCamelCase_   )
														lowerCAmelCase__          =	os.path.join(lowerCamelCase_					,    BART_VOCAB_FILES_NAMES['''vocab_file''']   )
														lowerCAmelCase__          =	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   __SCREAMING_SNAKE_CASE							(		self   )       ->							DPRQuestionEncoderTokenizer:
														return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,    '''dpr_tokenizer'''   )   )







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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[int]:
														shutil.rmtree(self.tmpdirname   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														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   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_=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(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    )
																					if init_retrieval:
																												retriever.init_retrieval(lowerCamelCase_   )
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							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=lowerCamelCase_   )
														lowerCAmelCase__          =	RagConfig(
														    retrieval_vector_size=self.retrieval_vector_size					,    question_encoder=DPRConfig().to_dict()					,    generator=BartConfig().to_dict()					,    )
														lowerCAmelCase__          =	ray.remote(lowerCamelCase_   )
														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(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    retrieval_workers=lowerCamelCase_					,    )
																					if init_retrieval:
																												retriever.init_retrieval()
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_=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(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    )
														else:
																					lowerCAmelCase__          =	RagPyTorchDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    index=CustomHFIndex(config.retrieval_vector_size					,    lowerCamelCase_   )					,    )
														if init_retrieval:
																					retriever.init_retrieval(lowerCamelCase_   )
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							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=lowerCamelCase_   )
														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(lowerCamelCase_   )
														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(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    retrieval_workers=lowerCamelCase_					,    index=CustomHFIndex.load_from_disk(
																					        vector_size=config.retrieval_vector_size					,    dataset_path=config.passages_path					,    index_path=config.index_path					,    )					,    )
														else:
																					lowerCAmelCase__          =	RagRayDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    retrieval_workers=lowerCamelCase_					,    index=CustomHFIndex(config.retrieval_vector_size					,    lowerCamelCase_   )					,    )
														if init_retrieval:
																					retriever.init_retrieval()
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							None:
														lowerCAmelCase__   ,					lowerCAmelCase__   ,					lowerCAmelCase__          =	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   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[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=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		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_custom_hf_index_pytorch_retriever(init_retrieval=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_					,    )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Dict:
														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=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_					,    )







							@require_ray
							def   __SCREAMING_SNAKE_CASE							(		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_ray_distributed_retriever(init_retrieval=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )
														ray.shutdown()







							@require_ray
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	np.array(
														    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]					,    dtype=np.floataa   )
														with self.assertRaises(lowerCamelCase_   ):
																					self.distributed_retriever_check(
																					    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_					,    )
														ray.shutdown()







							@require_ray
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[str]:
														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=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )
														ray.shutdown()