


"""simple docstring"""
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, PLBartTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


SCREAMING_SNAKE_CASE :    Dict									=				get_tests_dir("""fixtures/test_sentencepiece.model""")


if is_torch_available():
			from transformers.models.plbart.modeling_plbart import shift_tokens_right

SCREAMING_SNAKE_CASE :    Optional[int]									=				5_0003
SCREAMING_SNAKE_CASE :    str									=				5_0002



@require_sentencepiece
@require_tokenizers
class 			_UpperCAmelCase (							__snake_case,				unittest.TestCase				):






		'''simple docstring'''

		lowerCamelCase__				=PLBartTokenizer
		lowerCamelCase__				=None
		lowerCamelCase__				=False



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			super().setUp()

			# We have a SentencePiece fixture for testing
			__snake_case	:						List[Any]						=				PLBartTokenizer(a_				,       language_codes='''base'''				,       keep_accents=a_					)
			tokenizer.save_pretrained(self.tmpdirname					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				PLBartTokenizer(a_				,       language_codes='''base'''				,       keep_accents=a_					)

			__snake_case	:						List[Any]						=				tokenizer.tokenize('''This is a test'''					)
			self.assertListEqual(a_				,       ['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']					)

			self.assertListEqual(
			    tokenizer.convert_tokens_to_ids(a_					)				,       [value + tokenizer.fairseq_offset for value in [2_85, 46, 10, 1_70, 3_82]]				,       )

			__snake_case	:						Union[str, Any]						=				tokenizer.tokenize('''I was born in 92000, and this is falsé.'''					)
			self.assertListEqual(
			    a_				,       [
			        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''',
			        '''é''',
			        '''.''',
			    ]				,       )
			__snake_case	:						Dict						=				tokenizer.convert_tokens_to_ids(a_					)
			self.assertListEqual(
			    a_				,       [
			        value + tokenizer.fairseq_offset
			        for value in [8, 21, 84, 55, 24, 19, 7, 2, 6_02, 3_47, 3_47, 3_47, 3, 12, 66, 46, 72, 80, 6, 2, 4]
			    ]				,       )

			__snake_case	:						List[str]						=				tokenizer.convert_ids_to_tokens(a_					)
			self.assertListEqual(
			    a_				,       [
			        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>''',
			        '''.''',
			    ]				,       )

			__snake_case	:						List[Any]						=				tokenizer.vocab_size
			__snake_case	:						List[Any]						=				[tokenizer.convert_ids_to_tokens(a_					) for x in range(end - 4				,       a_					)]

			self.assertListEqual(a_				,       ['''__java__''', '''__python__''', '''__en_XX__''', '''<mask>''']					)

			__snake_case	:						int						=				'''java.lang.Exception, python.lang.Exception, javascript, php, ruby, go'''
			__snake_case	:						List[str]						=				tokenizer(a_					).input_ids
			self.assertEqual(
			    tokenizer.decode(a_				,       skip_special_tokens=a_				,       clean_up_tokenization_spaces=a_					)				,       a_				,       )



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						List[str]						=				PLBartTokenizer(a_				,       language_codes='''multi'''				,       keep_accents=a_					)

			__snake_case	:						int						=				tokenizer.tokenize('''This is a test'''					)
			self.assertListEqual(a_				,       ['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']					)

			self.assertListEqual(
			    tokenizer.convert_tokens_to_ids(a_					)				,       [value + tokenizer.fairseq_offset for value in [2_85, 46, 10, 1_70, 3_82]]				,       )

			__snake_case	:						str						=				tokenizer.tokenize('''I was born in 92000, and this is falsé.'''					)
			self.assertListEqual(
			    a_				,       [
			        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''',
			        '''é''',
			        '''.''',
			    ]				,       )
			__snake_case	:						Union[str, Any]						=				tokenizer.convert_tokens_to_ids(a_					)
			self.assertListEqual(
			    a_				,       [
			        value + tokenizer.fairseq_offset
			        for value in [8, 21, 84, 55, 24, 19, 7, 2, 6_02, 3_47, 3_47, 3_47, 3, 12, 66, 46, 72, 80, 6, 2, 4]
			    ]				,       )

			__snake_case	:						Any						=				tokenizer.convert_ids_to_tokens(a_					)
			self.assertListEqual(
			    a_				,       [
			        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>''',
			        '''.''',
			    ]				,       )
			__snake_case	:						List[Any]						=				tokenizer.vocab_size
			__snake_case	:						List[Any]						=				[tokenizer.convert_ids_to_tokens(a_					) for x in range(end - 7				,       a_					)]

			self.assertListEqual(
			    a_				,       ['''__java__''', '''__python__''', '''__en_XX__''', '''__javascript__''', '''__php__''', '''__ruby__''', '''__go__''']					)
			__snake_case	:						Dict						=				'''java.lang.Exception, python.lang.Exception, javascript, php, ruby, go'''
			__snake_case	:						int						=				tokenizer(a_					).input_ids
			self.assertEqual(
			    tokenizer.decode(a_				,       skip_special_tokens=a_				,       clean_up_tokenization_spaces=a_					)				,       a_				,       )



@require_torch
@require_sentencepiece
@require_tokenizers
class 			_UpperCAmelCase (							unittest.TestCase				):






		'''simple docstring'''

		lowerCamelCase__				='uclanlp/plbart-python-en_XX'
		lowerCamelCase__				=[
		    'def maximum(a,b,c):NEW_LINE_INDENTreturn max([a,b,c])',
		    'def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])',
		]
		lowerCamelCase__				=[
		    'Returns the maximum value of a b c.',
		    'Sums the values of a b c.',
		]
		lowerCamelCase__				=[
		    134,
		    5452,
		    33460,
		    33441,
		    33463,
		    33465,
		    33463,
		    33449,
		    988,
		    20,
		    33456,
		    19,
		    33456,
		    771,
		    39,
		    4258,
		    889,
		    3318,
		    33441,
		    33463,
		    33465,
		    33463,
		    33449,
		    2471,
		    2,
		    PYTHON_CODE,
		]



		@classmethod
		def 				SCREAMING_SNAKE_CASE  (cls					):





			'''simple docstring'''

			__snake_case	:						PLBartTokenizer						=				PLBartTokenizer.from_pretrained(
			    cls.checkpoint_name				,       language_codes='''base'''				,       src_lang='''python'''				,       tgt_lang='''en_XX'''					)
			__snake_case	:						Optional[int]						=				1
			return cls



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__java__''']				,       5_00_01					)
			self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__python__''']				,       5_00_02					)
			self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__en_XX__''']				,       5_00_03					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						int						=				self.tokenizer.batch_encode_plus(self.src_text					).input_ids[0]
			self.assertListEqual(self.expected_src_tokens				,       a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			self.assertIn(a_				,       self.tokenizer.all_special_ids					)
			__snake_case	:						Optional[int]						=				[EN_CODE, 90_37, 3_34_42, 57, 7_52, 1_53, 14, 56, 18, 9, 2]
			__snake_case	:						Tuple						=				self.tokenizer.decode(a_				,       skip_special_tokens=a_					)
			__snake_case	:						Optional[int]						=				self.tokenizer.decode(generated_ids[1:]				,       skip_special_tokens=a_					)
			self.assertEqual(a_				,       a_					)
			self.assertNotIn(self.tokenizer.eos_token				,       a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Tuple						=				['''def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])''' * 20]
			self.assertIsInstance(src_text[0]				,       a_					)
			__snake_case	:						int						=				10
			__snake_case	:						Optional[Any]						=				self.tokenizer(a_				,       max_length=a_				,       truncation=a_					).input_ids[0]
			self.assertEqual(ids[-2]				,       2					)
			self.assertEqual(ids[-1]				,       a_					)
			self.assertEqual(len(a_					)				,       a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['''<mask>''', '''__java__''']					)				,       [5_00_04, 5_00_01]					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Union[str, Any]						=				tempfile.mkdtemp()
			__snake_case	:						str						=				self.tokenizer.fairseq_tokens_to_ids
			self.tokenizer.save_pretrained(a_					)
			__snake_case	:						List[str]						=				PLBartTokenizer.from_pretrained(a_					)
			self.assertDictEqual(new_tok.fairseq_tokens_to_ids				,       a_					)



		@require_torch
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[int]						=				self.tokenizer(self.src_text				,       text_target=self.tgt_text				,       padding=a_				,       return_tensors='''pt'''					)
			__snake_case	:						Optional[int]						=				shift_tokens_right(batch['''labels''']				,       self.tokenizer.pad_token_id					)

			# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
			self.assertEqual(batch.input_ids[1][-2:].tolist()				,       [2, PYTHON_CODE]					)
			self.assertEqual(batch.decoder_input_ids[1][0]				,       a_					)
			self.assertEqual(batch.decoder_input_ids[1][-1]				,       2					)
			self.assertEqual(batch.labels[1][-2:].tolist()				,       [2, EN_CODE]					)



		@require_torch
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				self.tokenizer(
			    self.src_text				,       text_target=self.tgt_text				,       padding=a_				,       truncation=a_				,       max_length=len(self.expected_src_tokens					)				,       return_tensors='''pt'''				,       )
			__snake_case	:						List[str]						=				shift_tokens_right(batch['''labels''']				,       self.tokenizer.pad_token_id					)

			self.assertIsInstance(a_				,       a_					)

			self.assertEqual((2, 26)				,       batch.input_ids.shape					)
			self.assertEqual((2, 26)				,       batch.attention_mask.shape					)
			__snake_case	:						Optional[Any]						=				batch.input_ids.tolist()[0]
			self.assertListEqual(self.expected_src_tokens				,       a_					)
			self.assertEqual(2				,       batch.decoder_input_ids[0, -1]					)  # EOS
			# Test that special tokens are reset
			self.assertEqual(self.tokenizer.prefix_tokens				,       []					)
			self.assertEqual(self.tokenizer.suffix_tokens				,       [self.tokenizer.eos_token_id, PYTHON_CODE]					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						str						=				self.tokenizer(self.src_text				,       padding=a_				,       truncation=a_				,       max_length=3				,       return_tensors='''pt'''					)
			__snake_case	:						Tuple						=				self.tokenizer(
			    text_target=self.tgt_text				,       padding=a_				,       truncation=a_				,       max_length=10				,       return_tensors='''pt'''					)
			__snake_case	:						Optional[int]						=				targets['''input_ids''']
			__snake_case	:						Optional[Any]						=				shift_tokens_right(a_				,       self.tokenizer.pad_token_id					)

			self.assertEqual(batch.input_ids.shape[1]				,       3					)
			self.assertEqual(batch.decoder_input_ids.shape[1]				,       10					)



		@require_torch
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						List[str]						=				self.tokenizer._build_translation_inputs(
			    '''A test'''				,       return_tensors='''pt'''				,       src_lang='''en_XX'''				,       tgt_lang='''java'''					)

			self.assertEqual(
			    nested_simplify(a_					)				,       {
			        # A, test, EOS, en_XX
			        '''input_ids''': [[1_50, 2_42, 2, 5_00_03]],
			        '''attention_mask''': [[1, 1, 1, 1]],
			        # java
			        '''forced_bos_token_id''': 5_00_01,
			    }				,       )











"""simple docstring"""


import copy

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings


A_						=				r'''
    [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and
    can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.

    Args:
        title_sep (`str`, *optional*, defaults to  `" / "`):
            Separator inserted between the title and the text of the retrieved document when calling [`RagRetriever`].
        doc_sep (`str`, *optional*, defaults to  `" // "`):
            Separator inserted between the text of the retrieved document and the original input when calling
            [`RagRetriever`].
        n_docs (`int`, *optional*, defaults to 5):
            Number of documents to retrieve.
        max_combined_length (`int`, *optional*, defaults to 300):
            Max length of contextualized input returned by [`~RagRetriever.__call__`].
        retrieval_vector_size (`int`, *optional*, defaults to 768):
            Dimensionality of the document embeddings indexed by [`RagRetriever`].
        retrieval_batch_size (`int`, *optional*, defaults to 8):
            Retrieval batch size, defined as the number of queries issues concurrently to the faiss index encapsulated
            [`RagRetriever`].
        dataset (`str`, *optional*, defaults to `"wiki_dpr"`):
            A dataset identifier of the indexed dataset in HuggingFace Datasets (list all available datasets and ids
            using `datasets.list_datasets()`).
        dataset_split (`str`, *optional*, defaults to `"train"`)
            Which split of the `dataset` to load.
        index_name (`str`, *optional*, defaults to `"compressed"`)
            The index name of the index associated with the `dataset`. One can choose between `"legacy"`, `"exact"` and
            `"compressed"`.
        index_path (`str`, *optional*)
            The path to the serialized faiss index on disk.
        passages_path (`str`, *optional*):
            A path to text passages compatible with the faiss index. Required if using
            [`~models.rag.retrieval_rag.LegacyIndex`]
        use_dummy_dataset (`bool`, *optional*, defaults to `False`)
            Whether to load a "dummy" variant of the dataset specified by `dataset`.
        label_smoothing (`float`, *optional*, defaults to 0.0):
            Only relevant if `return_loss` is set to `True`. Controls the `epsilon` parameter value for label smoothing
            in the loss calculation. If set to 0, no label smoothing is performed.
        do_marginalize (`bool`, *optional*, defaults to `False`):
            If `True`, the logits are marginalized over all documents by making use of
            `torch.nn.functional.log_softmax`.
        reduce_loss (`bool`, *optional*, defaults to `False`):
            Whether or not to reduce the NLL loss using the `torch.Tensor.sum` operation.
        do_deduplication (`bool`, *optional*, defaults to `True`):
            Whether or not to deduplicate the generations from different context documents for a given input. Has to be
            set to `False` if used while training with distributed backend.
        exclude_bos_score (`bool`, *optional*, defaults to `False`):
            Whether or not to disregard the BOS token when computing the loss.
        output_retrieved(`bool`, *optional*, defaults to `False`):
            If set to `True`, `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and
            `context_attention_mask` are returned. See returned tensors for more detail.
        use_cache (`bool`, *optional*, defaults to `True`):
            Whether or not the model should return the last key/values attentions (not used by all models).
        forced_eos_token_id (`int`, *optional*):
            The id of the token to force as the last generated token when `max_length` is reached. Usually set to
            `eos_token_id`.
'''






@add_start_docstrings(__a )
class 					lowercase(     __a ):







 '''simple docstring'''





 lowercase__ 					=					"rag"
 lowercase__ 					=					True


 def __init__(				self:		Union[str, Any],      a_:		int=None,      a_:		Tuple=True,      a_:		Optional[int]=None,      a_:		List[str]=None,      a_:		int=None,      a_:		Optional[Any]=None,      a_:		List[str]=None,      a_:		Optional[Any]=" / ",      a_:		Tuple=" // ",      a_:		List[Any]=5,      a_:		Dict=300,      a_:		Tuple=768,      a_:		Optional[Any]=8,      a_:		int="wiki_dpr",      a_:		Any="train",      a_:		Optional[int]="compressed",      a_:		Optional[int]=None,      a_:		List[Any]=None,      a_:		Optional[Any]=False,      a_:		str=False,      a_:		Dict=0.0,      a_:		Union[str, Any]=True,      a_:		Union[str, Any]=False,      a_:		str=False,      a_:		List[str]=False,      a_:		Union[str, Any]=True,      a_:		Any=None,      **a_:		List[Any],      ):




   '''simple docstring'''



   super().__init__(
       bos_token_id=a_,      pad_token_id=a_,      eos_token_id=a_,      decoder_start_token_id=a_,      forced_eos_token_id=a_,      is_encoder_decoder=a_,      prefix=a_,      vocab_size=a_,      **a_,      )
   assert (
       "question_encoder" in kwargs and "generator" in kwargs
   ), "Config has to be initialized with question_encoder and generator config"
   _snake_case							:		Union[str, Any]								=  kwargs.pop("""question_encoder"""		)
   _snake_case							:		List[str]								=  question_encoder_config.pop("""model_type"""		)
   _snake_case							:		Union[str, Any]								=  kwargs.pop("""generator"""		)
   _snake_case							:		Any								=  decoder_config.pop("""model_type"""		)

   from ..auto.configuration_auto import AutoConfig

   _snake_case							:		Union[str, Any]								=  AutoConfig.for_model(a_,      **a_		)
   _snake_case							:		Optional[Any]								=  AutoConfig.for_model(a_,      **a_		)

   _snake_case							:		Any								=  reduce_loss
   _snake_case							:		Optional[int]								=  label_smoothing
   _snake_case							:		Dict								=  exclude_bos_score
   _snake_case							:		int								=  do_marginalize

   _snake_case							:		Optional[Any]								=  title_sep
   _snake_case							:		Any								=  doc_sep
   _snake_case							:		List[str]								=  n_docs
   _snake_case							:		Tuple								=  max_combined_length

   _snake_case							:		Optional[Any]								=  dataset
   _snake_case							:		Union[str, Any]								=  dataset_split
   _snake_case							:		Tuple								=  index_name

   _snake_case							:		Any								=  retrieval_vector_size
   _snake_case							:		Union[str, Any]								=  retrieval_batch_size
   _snake_case							:		str								=  passages_path
   _snake_case							:		Tuple								=  index_path
   _snake_case							:		List[Any]								=  use_dummy_dataset

   _snake_case							:		Optional[Any]								=  output_retrieved

   _snake_case							:		Tuple								=  do_deduplication

   _snake_case							:		Union[str, Any]								=  use_cache

   if self.forced_eos_token_id is None:
     _snake_case							:		Dict								=  getattr(self.generator,      """forced_eos_token_id""",      a_		)


 @classmethod
 def  UpperCamelCase_   (				cls:		Any,      a_:		PretrainedConfig,      a_:		PretrainedConfig,      **a_:		Optional[Any]		):




   '''simple docstring'''



   return cls(question_encoder=question_encoder_config.to_dict(),      generator=generator_config.to_dict(),      **a_		)


 def  UpperCamelCase_   (				self:		Tuple		):




   '''simple docstring'''



   _snake_case							:		Optional[int]								=  copy.deepcopy(self.__dict__		)
   _snake_case							:		List[str]								=  self.question_encoder.to_dict()
   _snake_case							:		Tuple								=  self.generator.to_dict()
   _snake_case							:		Dict								=  self.__class__.model_type
   return output


