





'''simple docstring'''


import unittest

from transformers import SPIECE_UNDERLINE, XLNetTokenizer, XLNetTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_lowerCamelCase :		List[Any] 			=				get_tests_dir('fixtures/test_sentencepiece.model')



@require_sentencepiece
@require_tokenizers
class 							__UpperCAmelCase     (				__lowerCAmelCase						,    unittest.TestCase  ):







		'''simple docstring'''







		__lowerCAmelCase        =			XLNetTokenizer
		__lowerCAmelCase        =			XLNetTokenizerFast
		__lowerCAmelCase        =			True
		__lowerCAmelCase        =			True







		def 				A   (self  :					Union[str, Any]					):
			super().setUp()

			# We have a SentencePiece fixture for testing
			A							     =					XLNetTokenizer(_lowerCAmelCase					,					keep_accents=_lowerCAmelCase					)
			tokenizer.sanitize_special_tokens()
			tokenizer.save_pretrained(self.tmpdirname					)







		def 				A   (self  :					Optional[Any]					):
			A							     =					'<s>'
			A							     =					1

			self.assertEqual(self.get_tokenizer()._convert_token_to_id(_lowerCAmelCase					)					,					_lowerCAmelCase					)
			self.assertEqual(self.get_tokenizer()._convert_id_to_token(_lowerCAmelCase					)					,					_lowerCAmelCase					)







		def 				A   (self  :					Tuple					):
			A							     =					list(self.get_tokenizer().get_vocab().keys()					)

			self.assertEqual(vocab_keys[0]					,					"""<unk>"""					)
			self.assertEqual(vocab_keys[1]					,					"""<s>"""					)
			self.assertEqual(vocab_keys[-1]					,					"""<eod>"""					)
			self.assertEqual(len(_lowerCAmelCase					)					,					1006					)







		def 				A   (self  :					int					):
			self.assertEqual(self.get_tokenizer().vocab_size					,					1000					)







		def 				A   (self  :					int					):
			A							     =					XLNetTokenizer(_lowerCAmelCase					,					keep_accents=_lowerCAmelCase					)

			A							     =					tokenizer.tokenize("""This is a test"""					)
			self.assertListEqual(_lowerCAmelCase					,					["""▁This""", """▁is""", """▁a""", """▁t""", """est"""]					)

			self.assertListEqual(tokenizer.convert_tokens_to_ids(_lowerCAmelCase					)					,					[285, 46, 10, 170, 382]					)

			A							     =					tokenizer.tokenize("""I was born in 92000, and this is falsé."""					)
			self.assertListEqual(
			    _lowerCAmelCase					,					[
			        SPIECE_UNDERLINE + """I""",
			        SPIECE_UNDERLINE + """was""",
			        SPIECE_UNDERLINE + """b""",
			        """or""",
			        """n""",
			        SPIECE_UNDERLINE + """in""",
			        SPIECE_UNDERLINE + """""",
			        """9""",
			        """2""",
			        """0""",
			        """0""",
			        """0""",
			        """,""",
			        SPIECE_UNDERLINE + """and""",
			        SPIECE_UNDERLINE + """this""",
			        SPIECE_UNDERLINE + """is""",
			        SPIECE_UNDERLINE + """f""",
			        """al""",
			        """s""",
			        """é""",
			        """.""",
			    ]					,					)
			A							     =					tokenizer.convert_tokens_to_ids(_lowerCAmelCase					)
			self.assertListEqual(_lowerCAmelCase					,					[8, 21, 84, 55, 24, 19, 7, 0, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 0, 4]					)

			A							     =					tokenizer.convert_ids_to_tokens(_lowerCAmelCase					)
			self.assertListEqual(
			    _lowerCAmelCase					,					[
			        SPIECE_UNDERLINE + """I""",
			        SPIECE_UNDERLINE + """was""",
			        SPIECE_UNDERLINE + """b""",
			        """or""",
			        """n""",
			        SPIECE_UNDERLINE + """in""",
			        SPIECE_UNDERLINE + """""",
			        """<unk>""",
			        """2""",
			        """0""",
			        """0""",
			        """0""",
			        """,""",
			        SPIECE_UNDERLINE + """and""",
			        SPIECE_UNDERLINE + """this""",
			        SPIECE_UNDERLINE + """is""",
			        SPIECE_UNDERLINE + """f""",
			        """al""",
			        """s""",
			        """<unk>""",
			        """.""",
			    ]					,					)







		def 				A   (self  :					List[Any]					):
			A							     =					XLNetTokenizer(_lowerCAmelCase					,					do_lower_case=_lowerCAmelCase					)
			A							     =					tokenizer.tokenize("""I was born in 92000, and this is falsé."""					)
			self.assertListEqual(
			    _lowerCAmelCase					,					[
			        SPIECE_UNDERLINE + """""",
			        """i""",
			        SPIECE_UNDERLINE + """was""",
			        SPIECE_UNDERLINE + """b""",
			        """or""",
			        """n""",
			        SPIECE_UNDERLINE + """in""",
			        SPIECE_UNDERLINE + """""",
			        """9""",
			        """2""",
			        """0""",
			        """0""",
			        """0""",
			        """,""",
			        SPIECE_UNDERLINE + """and""",
			        SPIECE_UNDERLINE + """this""",
			        SPIECE_UNDERLINE + """is""",
			        SPIECE_UNDERLINE + """f""",
			        """al""",
			        """se""",
			        """.""",
			    ]					,					)
			self.assertListEqual(tokenizer.tokenize("""H\u00E9llo"""					)					,					["""▁he""", """ll""", """o"""]					)







		def 				A   (self  :					List[str]					):
			A							     =					XLNetTokenizer(_lowerCAmelCase					,					do_lower_case=_lowerCAmelCase					)
			A							     =					tokenizer.tokenize("""I was born in 92000, and this is falsé."""					)
			self.assertListEqual(
			    _lowerCAmelCase					,					[
			        SPIECE_UNDERLINE + """I""",
			        SPIECE_UNDERLINE + """was""",
			        SPIECE_UNDERLINE + """b""",
			        """or""",
			        """n""",
			        SPIECE_UNDERLINE + """in""",
			        SPIECE_UNDERLINE + """""",
			        """9""",
			        """2""",
			        """0""",
			        """0""",
			        """0""",
			        """,""",
			        SPIECE_UNDERLINE + """and""",
			        SPIECE_UNDERLINE + """this""",
			        SPIECE_UNDERLINE + """is""",
			        SPIECE_UNDERLINE + """f""",
			        """al""",
			        """se""",
			        """.""",
			    ]					,					)







		@slow
		def 				A   (self  :					Dict					):
			A							     =					XLNetTokenizer.from_pretrained("""xlnet-base-cased"""					)

			A							     =					tokenizer.encode("""sequence builders"""					,					add_special_tokens=_lowerCAmelCase					)
			A							     =					tokenizer.encode("""multi-sequence build"""					,					add_special_tokens=_lowerCAmelCase					)

			A							     =					tokenizer.build_inputs_with_special_tokens(_lowerCAmelCase					)
			A							     =					tokenizer.build_inputs_with_special_tokens(_lowerCAmelCase					,					_lowerCAmelCase					)

			assert encoded_sentence == text + [4, 3]
			assert encoded_pair == text + [4] + text_a + [4, 3]







		@slow
		def 				A   (self  :					Tuple					):
			A							     =					{'input_ids': [[17, 2_1442, 270, 17, 10, 1_4645, 318, 34, 17, 4546, 3145, 787, 13, 7752, 2_2018, 23, 21, 17, 4546, 3145, 787, 13, 3352, 1_4431, 13, 5500, 11, 1176, 580, 13, 1_6819, 4797, 23, 17, 10, 1_7135, 658, 19, 457, 7932, 13, 184, 19, 3154, 1_7135, 6468, 19, 1404, 1_2269, 19, 4229, 5356, 1_6264, 46, 19, 17, 2_0545, 1_0395, 9, 9, 9, 11, 28, 6421, 9531, 2_0729, 17, 10, 353, 1_7022, 11, 21, 6421, 9531, 1_6949, 17, 10, 1_1509, 753, 11, 33, 95, 2421, 7385, 956, 1_4431, 2626, 25, 842, 7385, 4836, 21, 1429, 2272, 9855, 3120, 161, 2_4738, 19, 1_3203, 658, 218, 787, 21, 430, 1_8482, 847, 2637, 9, 4, 3], [5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 322, 2_2178, 27, 1064, 22, 956, 13, 1_1101, 1429, 5854, 2_4313, 1_8953, 40, 422, 2_4366, 68, 1758, 37, 1_0483, 1_4257, 31, 207, 263, 21, 203, 3773, 25, 71, 9735, 9, 4, 3], [5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 32, 2049, 3442, 17, 1_3894, 3380, 23, 95, 18, 1_7634, 2288, 9, 4, 3]], 'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2], [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2], [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}  # noqa: E501
			# fmt: on

			self.tokenizer_integration_test_util(
			    expected_encoding=_lowerCAmelCase					,					model_name="""xlnet-base-cased"""					,					revision="""c841166438c31ec7ca9a106dee7bb312b73ae511"""					,					)


"""simple docstring"""
import logging
import random

import ray

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


lowercase__					       =	logging.getLogger(__name__)




class    __snake_case       :



      def __init__(						self)	->			Optional[int]:





        '''simple docstring'''





        a__:     Optional[Any]										=	False



      def   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase    ,			lowercase    ,			lowercase)	->			str:





        '''simple docstring'''





        if not self.initialized:
          a__:     Optional[int]										=	RagRetriever(
              lowercase    ,			question_encoder_tokenizer=lowercase    ,			generator_tokenizer=lowercase    ,			index=lowercase    ,			init_retrieval=lowercase    ,			)
          a__:     Optional[int]										=	True



      def   lowerCamelCase_   (						self)	->			int:





        '''simple docstring'''





        self.retriever.index.init_index()



      def   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase)	->			Union[str, Any]:





        '''simple docstring'''





        a__							,    a__:     str										=	self.retriever._main_retrieve(lowercase    ,			lowercase)
        return doc_ids, retrieved_doc_embeds







class    __snake_case       (     __lowerCAmelCase				):



      def __init__(						self    ,			lowercase    ,			lowercase    ,			lowercase    ,			lowercase    ,			lowercase=None)	->			int:





        '''simple docstring'''





        if index is not None and index.is_initialized() and len(lowercase) > 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__(
            lowercase    ,			question_encoder_tokenizer=lowercase    ,			generator_tokenizer=lowercase    ,			index=lowercase    ,			init_retrieval=lowercase    ,			)
        a__:     Any										=	retrieval_workers
        if len(self.retrieval_workers) > 0:
          ray.get(
              [
                  worker.create_rag_retriever.remote(lowercase    ,			lowercase    ,			lowercase    ,			lowercase)
                  for worker in self.retrieval_workers
              ])



      def   lowerCamelCase_   (						self)	->			Any:





        '''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   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase)	->			Union[str, Any]:





        '''simple docstring'''





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



      @classmethod
      def   lowerCamelCase_   (						cls    ,			lowercase    ,			lowercase=None    ,			**lowercase)	->			Tuple:





        '''simple docstring'''





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



      @classmethod
      def   lowerCamelCase_   (						cls    ,			lowercase    ,			lowercase    ,			lowercase=None    ,			**lowercase)	->			Union[str, Any]:





        '''simple docstring'''





        a__:     Optional[int]										=	kwargs.pop('config'    ,			lowercase) or RagConfig.from_pretrained(lowercase    ,			**lowercase)
        a__:     Union[str, Any]										=	RagTokenizer.from_pretrained(lowercase    ,			config=lowercase)
        a__:     int										=	rag_tokenizer.question_encoder
        a__:     Any										=	rag_tokenizer.generator
        if indexed_dataset is not None:
          a__:     List[Any]										=	'custom'
          a__:     Optional[Any]										=	CustomHFIndex(config.retrieval_vector_size    ,			lowercase)
        else:
          a__:     Dict										=	cls._build_index(lowercase)
        return cls(
            lowercase    ,			question_encoder_tokenizer=lowercase    ,			generator_tokenizer=lowercase    ,			retrieval_workers=lowercase    ,			index=lowercase    ,			)


