




'''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     UpperCamelCase_     (a__,      unittest.TestCase     ):







						"""simple docstring"""


						_lowerCAmelCase    						=					GPTNeoXJapaneseTokenizer
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					{'do_clean_text': False, 'add_prefix_space': False}







						def       _a    (     self      :       Union[str, Any]    ):





											"""simple docstring"""







											super().setUp()

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

											A_ :     Dict =							os.path.join(self.tmpdirname							, VOCAB_FILES_NAMES['''vocab_file''']    )
											A_ :     Optional[Any] =							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       _a    (     self      :       str							, **_lowerCamelCase      :       Any    ):





											"""simple docstring"""







											kwargs.update(self.special_tokens_map    )
											return GPTNeoXJapaneseTokenizer.from_pretrained(self.tmpdirname							, **_lowerCamelCase    )







						def       _a    (     self      :       Dict							, _lowerCamelCase      :       Dict    ):





											"""simple docstring"""







											A_ :     Optional[int] =							'''こんにちは、世界。 \nこんばんは、㔺界。😀'''
											A_ :     Union[str, Any] =							'''こんにちは、世界。 \nこんばんは、世界。😀'''
											return input_text, output_text







						def       _a    (     self      :       Tuple							, _lowerCamelCase      :       List[str]    ):





											"""simple docstring"""







											A_ ,A_ :     List[str] =							self.get_input_output_texts(_lowerCamelCase    )
											A_ :     str =							tokenizer.encode(_lowerCamelCase							, add_special_tokens=_lowerCamelCase    )
											A_ :     Optional[Any] =							tokenizer.decode(_lowerCamelCase							, clean_up_tokenization_spaces=_lowerCamelCase    )
											return text, ids







						def       _a    (     self      :       Union[str, Any]    ):





											"""simple docstring"""







											pass  # TODO add if relevant







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											pass  # TODO add if relevant







						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											pass  # TODO add if relevant







						def       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ :     List[str] =							self.get_tokenizer()

											# Testing tokenization
											A_ :     int =							'''こんにちは、世界。　こんばんは、㔺界。'''
											A_ :     Tuple =							['''こん''', '''にちは''', '''、''', '''世界''', '''。''', '''<SP>''', '''こん''', '''ばんは''', '''、''', '''㔺界''', '''。''']
											A_ :     List[str] =							tokenizer.tokenize(_lowerCamelCase    )
											self.assertListEqual(_lowerCamelCase							, _lowerCamelCase    )

											# Testing conversion to ids without special tokens
											A_ :     Dict =							[0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6]
											A_ :     str =							tokenizer.convert_tokens_to_ids(_lowerCamelCase    )
											self.assertListEqual(_lowerCamelCase							, _lowerCamelCase    )

											# Testing conversion to ids with special tokens
											A_ :     Optional[Any] =							tokens + [tokenizer.unk_token]
											A_ :     Any =							[0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6, 19]
											A_ :     Tuple =							tokenizer.convert_tokens_to_ids(_lowerCamelCase    )
											self.assertListEqual(_lowerCamelCase							, _lowerCamelCase    )







						@slow
						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							self.tokenizer_class.from_pretrained('''abeja/gpt-neox-japanese-2.7b'''    )

											A_ :     str =							tokenizer.encode('''ありがとう。'''							, add_special_tokens=_lowerCamelCase    )
											A_ :     Optional[Any] =							tokenizer.encode('''どういたしまして。'''							, add_special_tokens=_lowerCamelCase    )

											A_ :     str =							tokenizer.build_inputs_with_special_tokens(_lowerCamelCase    )
											A_ :     Any =							tokenizer.build_inputs_with_special_tokens(_lowerCamelCase							, _lowerCamelCase    )

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







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											pass







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											pass









'''simple docstring'''




import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


snake_case__          =						logging.getLogger(__name__)






class     UpperCamelCase_     :







						"""simple docstring"""


						def __init__(     self      :       Tuple    ):





											"""simple docstring"""







											A_ :     Optional[int] =							False







						def       _a    (     self      :       Any							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       str							, _lowerCamelCase      :       int    ):





											"""simple docstring"""







											if not self.initialized:
																A_ :     Optional[Any] =							RagRetriever(
																    _lowerCamelCase							, question_encoder_tokenizer=_lowerCamelCase							, generator_tokenizer=_lowerCamelCase							, index=_lowerCamelCase							, init_retrieval=_lowerCamelCase							, )
																A_ :     Optional[Any] =							True







						def       _a    (     self      :       int    ):





											"""simple docstring"""







											self.retriever.index.init_index()







						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											del self.retriever
											A_ :     List[Any] =							False







						def       _a    (     self      :       Optional[int]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ ,A_ :     Dict =							self.retriever._main_retrieve(_lowerCamelCase							, _lowerCamelCase    )
											A_ :     int =							self.retriever.index.get_doc_dicts(_lowerCamelCase    )
											return doc_ids, retrieved_doc_embeds, doc_dicts






class     UpperCamelCase_     (a__     ):







						"""simple docstring"""


						def __init__(     self      :       int							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Tuple=None    ):





											"""simple docstring"""







											if index is not None and index.is_initialized() and len(_lowerCamelCase    ) > 0:
																raise ValueError(
																    '''When using Ray for distributed fine-tuning, '''
																    '''you\'ll need to provide the paths instead, '''
																    '''as the dataset and the index are loaded '''
																    '''separately. More info in examples/rag/use_own_knowledge_dataset.py '''    )

											super().__init__(
											    _lowerCamelCase							, question_encoder_tokenizer=_lowerCamelCase							, generator_tokenizer=_lowerCamelCase							, index=_lowerCamelCase							, init_retrieval=_lowerCamelCase							, )

											A_ :     Tuple =							retrieval_workers
											A_ :     Any =							question_encoder_tokenizer
											A_ :     Dict =							generator_tokenizer
											if len(self.retrieval_workers    ) > 0:
																ray.get(
																    [
																        worker.create_rag_retriever.remote(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    )
																        for worker in self.retrieval_workers
																    ]    )







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											logger.info('''initializing retrieval'''    )

											if len(self.retrieval_workers    ) > 0:
																ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]    )
											else:
																# Non-distributed training. Load index into this same process.
																self.index.init_index()







						def       _a    (     self      :       Dict							, _lowerCamelCase      :       Tuple							, _lowerCamelCase      :       int    ):





											"""simple docstring"""







											if len(self.retrieval_workers    ) > 0:
																# Select a random retrieval actor.
																A_ :     Optional[int] =							self.retrieval_workers[random.randint(0							, len(self.retrieval_workers    ) - 1    )]
																A_ ,A_ ,A_ :     Union[str, Any] =							ray.get(
																    random_worker.retrieve.remote(_lowerCamelCase							, _lowerCamelCase    )    )
											else:
																A_ ,A_ :     List[str] =							self._main_retrieve(_lowerCamelCase							, _lowerCamelCase    )
																A_ :     Union[str, Any] =							self.index.get_doc_dicts(_lowerCamelCase    )
											return retrieved_doc_embeds, doc_ids, doc_dicts







						@classmethod
						def       _a    (     cls      :       List[str]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Optional[Any]=None							, **_lowerCamelCase      :       str    ):





											"""simple docstring"""







											return super(_lowerCamelCase							, cls    ).get_tokenizers(_lowerCamelCase							, _lowerCamelCase							, **_lowerCamelCase    )







						@classmethod
						def       _a    (     cls      :       Optional[Any]							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       List[Any]							, _lowerCamelCase      :       str=None							, **_lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







											A_ :     Tuple =							kwargs.pop('''config'''							, _lowerCamelCase    ) or RagConfig.from_pretrained(_lowerCamelCase							, **_lowerCamelCase    )
											A_ :     List[Any] =							RagTokenizer.from_pretrained(_lowerCamelCase							, config=_lowerCamelCase    )
											A_ :     Union[str, Any] =							rag_tokenizer.question_encoder
											A_ :     Dict =							rag_tokenizer.generator

											if indexed_dataset is not None:
																A_ :     Any =							'''custom'''
																A_ :     int =							CustomHFIndex(config.retrieval_vector_size							, _lowerCamelCase    )
											else:
																A_ :     Optional[int] =							cls._build_index(_lowerCamelCase    )

											return cls(
											    _lowerCamelCase							, question_encoder_tokenizer=_lowerCamelCase							, generator_tokenizer=_lowerCamelCase							, retrieval_workers=_lowerCamelCase							, index=_lowerCamelCase							, )







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											logger.info('''re-loading the new dataset with embeddings'''    )
											# access from the training loop

											ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]    )

											# build the index object again
											A_ :     Dict =							self._build_index(self.config    )

											ray.get(
											    [
											        worker.create_rag_retriever.remote(
											            self.config							, self.question_encoder_tokenizer							, self.generator_tokenizer							, _lowerCamelCase    )
											        for worker in self.retrieval_workers
											    ]    )



