


"""simple docstring"""






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

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
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_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
							import faiss





@require_faiss
class 	lowerCamelCase__  (	__a							):







					"""simple docstring"""


					def lowerCamelCase__				(			self       :					List[Any]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Union[str, Any] 				=   tempfile.mkdtemp()
												__UpperCAmelCase  :		Dict 				=   8

												# DPR tok
												__UpperCAmelCase  :		Optional[int] 				=   [
												    '''[UNK]''',
												    '''[CLS]''',
												    '''[SEP]''',
												    '''[PAD]''',
												    '''[MASK]''',
												    '''want''',
												    '''##want''',
												    '''##ed''',
												    '''wa''',
												    '''un''',
												    '''runn''',
												    '''##ing''',
												    ''',''',
												    '''low''',
												    '''lowest''',
												]
												__UpperCAmelCase  :		int 				=   os.path.join(self.tmpdirname					,			"""dpr_tokenizer"""    )
												os.makedirs(_lowercase					,			exist_ok=_lowercase    )
												__UpperCAmelCase  :		Optional[Any] 				=   os.path.join(_lowercase					,			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
												__UpperCAmelCase  :		Tuple 				=   [
												    '''l''',
												    '''o''',
												    '''w''',
												    '''e''',
												    '''r''',
												    '''s''',
												    '''t''',
												    '''i''',
												    '''d''',
												    '''n''',
												    '''\u0120''',
												    '''\u0120l''',
												    '''\u0120n''',
												    '''\u0120lo''',
												    '''\u0120low''',
												    '''er''',
												    '''\u0120lowest''',
												    '''\u0120newer''',
												    '''\u0120wider''',
												    '''<unk>''',
												]
												__UpperCAmelCase  :		int 				=   dict(zip(_lowercase					,			range(len(_lowercase    )    )    )    )
												__UpperCAmelCase  :		Tuple 				=   ['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
												__UpperCAmelCase  :		List[str] 				=   {'''unk_token''': '''<unk>'''}

												__UpperCAmelCase  :		Dict 				=   os.path.join(self.tmpdirname					,			"""bart_tokenizer"""    )
												os.makedirs(_lowercase					,			exist_ok=_lowercase    )
												__UpperCAmelCase  :		Dict 				=   os.path.join(_lowercase					,			BART_VOCAB_FILES_NAMES["""vocab_file"""]    )
												__UpperCAmelCase  :		Optional[Any] 				=   os.path.join(_lowercase					,			BART_VOCAB_FILES_NAMES["""merges_file"""]    )
												with open(self.vocab_file					,			"""w"""					,			encoding="""utf-8"""    ) as fp:
																			fp.write(json.dumps(_lowercase    ) + """\n"""    )
												with open(self.merges_file					,			"""w"""					,			encoding="""utf-8"""    ) as fp:
																			fp.write("""\n""".join(_lowercase    )    )






					def lowerCamelCase__				(			self       :					Union[str, Any]    ):
												'''simple docstring'''






												return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,			"""dpr_tokenizer"""    )    )






					def lowerCamelCase__				(			self       :					List[Any]    ):
												'''simple docstring'''






												return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,			"""dpr_tokenizer"""    )    )






					def lowerCamelCase__				(			self       :					Any    ):
												'''simple docstring'''






												return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname					,			"""bart_tokenizer"""    )    )






					def lowerCamelCase__				(			self       :					Optional[int]    ):
												'''simple docstring'''






												shutil.rmtree(self.tmpdirname    )






					def lowerCamelCase__				(			self       :					List[str]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		List[str] 				=   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       :					Any    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Union[str, Any] 				=   self.get_dummy_dataset()
												__UpperCAmelCase  :		List[Any] 				=   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:
																			__UpperCAmelCase  :		List[Any] 				=   dataset
																			__UpperCAmelCase  :		str 				=   RagRetriever(
																			    _lowercase					,			question_encoder_tokenizer=self.get_dpr_tokenizer()					,			generator_tokenizer=self.get_bart_tokenizer()					,			)
												return retriever






					def lowerCamelCase__				(			self       :					str					,			UpperCamelCase       :					bool    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Optional[Any] 				=   self.get_dummy_dataset()
												__UpperCAmelCase  :		Optional[int] 				=   RagConfig(
												    retrieval_vector_size=self.retrieval_vector_size					,			question_encoder=DPRConfig().to_dict()					,			generator=BartConfig().to_dict()					,			index_name="""custom"""					,			)
												if from_disk:
																			__UpperCAmelCase  :		int 				=   os.path.join(self.tmpdirname					,			"""dataset"""    )
																			__UpperCAmelCase  :		Union[str, Any] 				=   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
																			__UpperCAmelCase  :		Dict 				=   RagRetriever(
																			    _lowercase					,			question_encoder_tokenizer=self.get_dpr_tokenizer()					,			generator_tokenizer=self.get_bart_tokenizer()					,			)
												else:
																			__UpperCAmelCase  :		List[str] 				=   RagRetriever(
																			    _lowercase					,			question_encoder_tokenizer=self.get_dpr_tokenizer()					,			generator_tokenizer=self.get_bart_tokenizer()					,			index=CustomHFIndex(config.retrieval_vector_size					,			_lowercase    )					,			)
												return retriever






					def lowerCamelCase__				(			self       :					List[str]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		str 				=   Dataset.from_dict(
												    {
												        """id""": ["""0""", """1"""],
												        """text""": ["""foo""", """bar"""],
												        """title""": ["""Foo""", """Bar"""],
												        """embeddings""": [np.ones(self.retrieval_vector_size + 1    ), 2 * np.ones(self.retrieval_vector_size + 1    )],
												    }    )
												dataset.add_faiss_index("""embeddings"""					,			string_factory="""Flat"""					,			metric_type=faiss.METRIC_INNER_PRODUCT    )

												__UpperCAmelCase  :		Dict 				=   os.path.join(self.tmpdirname					,			"""hf_bert_base.hnswSQ8_correct_phi_128.c_index"""    )
												dataset.save_faiss_index("""embeddings"""					,			index_file_name + """.index.dpr"""    )
												pickle.dump(dataset["""id"""]					,			open(index_file_name + """.index_meta.dpr"""					,			"""wb"""    )    )

												__UpperCAmelCase  :		List[Any] 				=   os.path.join(self.tmpdirname					,			"""psgs_w100.tsv.pkl"""    )
												__UpperCAmelCase  :		List[str] 				=   {sample['''id''']: [sample['''text'''], sample['''title''']] for sample in dataset}
												pickle.dump(_lowercase					,			open(_lowercase					,			"""wb"""    )    )

												__UpperCAmelCase  :		List[Any] 				=   RagConfig(
												    retrieval_vector_size=self.retrieval_vector_size					,			question_encoder=DPRConfig().to_dict()					,			generator=BartConfig().to_dict()					,			index_name="""legacy"""					,			index_path=self.tmpdirname					,			)
												__UpperCAmelCase  :		Any 				=   RagRetriever(
												    _lowercase					,			question_encoder_tokenizer=self.get_dpr_tokenizer()					,			generator_tokenizer=self.get_bart_tokenizer()    )
												return retriever






					def lowerCamelCase__				(			self       :					List[str]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		List[str] 				=   1
												__UpperCAmelCase  :		List[str] 				=   self.get_dummy_canonical_hf_index_retriever()
												__UpperCAmelCase  :		Optional[int] 				=   np.array(
												    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
												__UpperCAmelCase  :		Optional[int] 				=   retriever.retrieve(_lowercase					,			n_docs=_lowercase    )
												self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)    )
												self.assertEqual(len(_lowercase    )					,			2    )
												self.assertEqual(sorted(doc_dicts[0]    )					,			["""embeddings""", """id""", """text""", """title"""]    )
												self.assertEqual(len(doc_dicts[0]["""id"""]    )					,			_lowercase    )
												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    ):
												'''simple docstring'''






												__UpperCAmelCase  :		int 				=   self.get_dummy_canonical_hf_index_retriever()
												with tempfile.TemporaryDirectory() as tmp_dirname:
																			with patch("""transformers.models.rag.retrieval_rag.load_dataset"""    ) as mock_load_dataset:
																										__UpperCAmelCase  :		Optional[int] 				=   self.get_dummy_dataset()
																										retriever.save_pretrained(_lowercase    )
																										__UpperCAmelCase  :		Any 				=   RagRetriever.from_pretrained(_lowercase    )
																										self.assertIsInstance(_lowercase					,			_lowercase    )
																										__UpperCAmelCase  :		int 				=   np.array(
																										    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
																										__UpperCAmelCase  :		Tuple 				=   retriever.retrieve(_lowercase					,			n_docs=1    )
																										self.assertTrue(out is not None    )






					def lowerCamelCase__				(			self       :					Optional[int]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		str 				=   1
												__UpperCAmelCase  :		str 				=   self.get_dummy_custom_hf_index_retriever(from_disk=_lowercase    )
												__UpperCAmelCase  :		List[Any] 				=   np.array(
												    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
												__UpperCAmelCase  :		Any 				=   retriever.retrieve(_lowercase					,			n_docs=_lowercase    )
												self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)    )
												self.assertEqual(len(_lowercase    )					,			2    )
												self.assertEqual(sorted(doc_dicts[0]    )					,			["""embeddings""", """id""", """text""", """title"""]    )
												self.assertEqual(len(doc_dicts[0]["""id"""]    )					,			_lowercase    )
												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       :					str    ):
												'''simple docstring'''






												__UpperCAmelCase  :		str 				=   self.get_dummy_custom_hf_index_retriever(from_disk=_lowercase    )
												with tempfile.TemporaryDirectory() as tmp_dirname:
																			retriever.save_pretrained(_lowercase    )
																			__UpperCAmelCase  :		Optional[int] 				=   RagRetriever.from_pretrained(_lowercase    )
																			self.assertIsInstance(_lowercase					,			_lowercase    )
																			__UpperCAmelCase  :		Union[str, Any] 				=   np.array(
																			    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
																			__UpperCAmelCase  :		Any 				=   retriever.retrieve(_lowercase					,			n_docs=1    )
																			self.assertTrue(out is not None    )






					def lowerCamelCase__				(			self       :					Union[str, Any]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		List[Any] 				=   1
												__UpperCAmelCase  :		Optional[int] 				=   self.get_dummy_custom_hf_index_retriever(from_disk=_lowercase    )
												__UpperCAmelCase  :		Optional[int] 				=   np.array(
												    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
												__UpperCAmelCase  :		str 				=   retriever.retrieve(_lowercase					,			n_docs=_lowercase    )
												self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)    )
												self.assertEqual(len(_lowercase    )					,			2    )
												self.assertEqual(sorted(doc_dicts[0]    )					,			["""embeddings""", """id""", """text""", """title"""]    )
												self.assertEqual(len(doc_dicts[0]["""id"""]    )					,			_lowercase    )
												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       :					List[str]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Dict 				=   self.get_dummy_custom_hf_index_retriever(from_disk=_lowercase    )
												with tempfile.TemporaryDirectory() as tmp_dirname:
																			retriever.save_pretrained(_lowercase    )
																			__UpperCAmelCase  :		Union[str, Any] 				=   RagRetriever.from_pretrained(_lowercase    )
																			self.assertIsInstance(_lowercase					,			_lowercase    )
																			__UpperCAmelCase  :		str 				=   np.array(
																			    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
																			__UpperCAmelCase  :		List[str] 				=   retriever.retrieve(_lowercase					,			n_docs=1    )
																			self.assertTrue(out is not None    )






					def lowerCamelCase__				(			self       :					Any    ):
												'''simple docstring'''






												__UpperCAmelCase  :		List[str] 				=   1
												__UpperCAmelCase  :		Any 				=   self.get_dummy_legacy_index_retriever()
												__UpperCAmelCase  :		Union[str, Any] 				=   np.array(
												    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
												__UpperCAmelCase  :		Union[str, Any] 				=   retriever.retrieve(_lowercase					,			n_docs=_lowercase    )
												self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)    )
												self.assertEqual(len(_lowercase    )					,			2    )
												self.assertEqual(sorted(doc_dicts[0]    )					,			["""text""", """title"""]    )
												self.assertEqual(len(doc_dicts[0]["""text"""]    )					,			_lowercase    )
												self.assertEqual(doc_dicts[0]["""text"""][0]					,			"""bar"""    )  # max inner product is reached with second doc
												self.assertEqual(doc_dicts[1]["""text"""][0]					,			"""foo"""    )  # max inner product is reached with first doc
												self.assertListEqual(doc_ids.tolist()					,			[[1], [0]]    )






					def lowerCamelCase__				(			self       :					int    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Dict 				=   self.get_dummy_legacy_index_retriever()
												with tempfile.TemporaryDirectory() as tmp_dirname:
																			retriever.save_pretrained(_lowercase    )
																			__UpperCAmelCase  :		Dict 				=   RagRetriever.from_pretrained(_lowercase    )
																			self.assertIsInstance(_lowercase					,			_lowercase    )
																			__UpperCAmelCase  :		Any 				=   np.array(
																			    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
																			__UpperCAmelCase  :		Optional[int] 				=   retriever.retrieve(_lowercase					,			n_docs=1    )
																			self.assertTrue(out is not None    )






					@require_torch
					@require_tokenizers
					@require_sentencepiece
					def lowerCamelCase__				(			self       :					str    ):
												'''simple docstring'''






												import torch

												__UpperCAmelCase  :		List[str] 				=   1
												__UpperCAmelCase  :		Dict 				=   self.get_dummy_canonical_hf_index_retriever()
												__UpperCAmelCase  :		Any 				=   [[5, 7], [10, 11]]
												__UpperCAmelCase  :		int 				=   np.array(
												    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
												__UpperCAmelCase  :		Optional[int] 				=   retriever(_lowercase					,			_lowercase					,			prefix=retriever.config.generator.prefix					,			n_docs=_lowercase    )
												__UpperCAmelCase  :		List[Any] 				=   (
												    out['''context_input_ids'''],
												    out['''context_attention_mask'''],
												    out['''retrieved_doc_embeds'''],
												)
												self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)    )
												self.assertIsInstance(_lowercase					,			_lowercase    )
												self.assertIsInstance(_lowercase					,			_lowercase    )
												self.assertIsInstance(_lowercase					,			np.ndarray    )

												__UpperCAmelCase  :		Union[str, Any] 				=   retriever(
												    _lowercase					,			_lowercase					,			prefix=retriever.config.generator.prefix					,			n_docs=_lowercase					,			return_tensors="""pt"""					,			)
												__UpperCAmelCase  :		Tuple 				=   (  # noqa: F841
												    out['''context_input_ids'''],
												    out['''context_attention_mask'''],
												    out['''retrieved_doc_embeds'''],
												    out['''doc_ids'''],
												)
												self.assertEqual(retrieved_doc_embeds.shape					,			(2, n_docs, self.retrieval_vector_size)    )
												self.assertIsInstance(_lowercase					,			torch.Tensor    )
												self.assertIsInstance(_lowercase					,			torch.Tensor    )
												self.assertIsInstance(_lowercase					,			torch.Tensor    )






					@require_torch
					@require_tokenizers
					@require_sentencepiece
					def lowerCamelCase__				(			self       :					Optional[int]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Optional[Any] 				=   self.get_dpr_ctx_encoder_tokenizer()
												__UpperCAmelCase  :		List[str] 				=   1
												__UpperCAmelCase  :		Union[str, Any] 				=   self.get_dummy_custom_hf_index_retriever(from_disk=_lowercase    )
												retriever.set_ctx_encoder_tokenizer(_lowercase    )

												__UpperCAmelCase  :		str 				=   [[5, 7], [10, 11]]
												__UpperCAmelCase  :		str 				=   np.array(
												    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]					,			dtype=np.floataa    )
												__UpperCAmelCase  :		Optional[int] 				=   retriever(_lowercase					,			_lowercase					,			prefix=retriever.config.generator.prefix					,			n_docs=_lowercase    )

												self.assertEqual(
												    len(_lowercase    )					,			6    )  # check whether the retriever output consist of 6 attributes including tokenized docs
												self.assertEqual(
												    all(k in out for k in ("""tokenized_doc_ids""", """tokenized_doc_attention_mask""")    )					,			_lowercase    )  # check for doc token related keys in dictionary.

from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TFLEDForConditionalGeneration, TFLEDModel




@require_tf
class 			__magic_name__	:





   """simple docstring"""




   lowerCAmelCase   :     Optional[int]	      = LEDConfig
   lowerCAmelCase   :     List[str]	      = {}
   lowerCAmelCase   :     str	      = '''gelu'''



   def __init__(						self       :	int		,	_lowercase       :	Tuple		,	_lowercase       :	Dict=13		,	_lowercase       :	Union[str, Any]=7		,	_lowercase       :	Optional[int]=True		,	_lowercase       :	Any=False		,	_lowercase       :	Optional[int]=99		,	_lowercase       :	Dict=32		,	_lowercase       :	Optional[Any]=2		,	_lowercase       :	Union[str, Any]=4		,	_lowercase       :	List[str]=37		,	_lowercase       :	Optional[Any]=0.1		,	_lowercase       :	List[Any]=0.1		,	_lowercase       :	Union[str, Any]=20		,	_lowercase       :	Tuple=2		,	_lowercase       :	str=1		,	_lowercase       :	Optional[int]=0		,	_lowercase       :	List[str]=4		,	):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			parent
    _UpperCamelCase:  Union[str, Any]    						=			batch_size
    _UpperCamelCase:  Optional[Any]    						=			seq_length
    _UpperCamelCase:  str    						=			is_training
    _UpperCamelCase:  Dict    						=			use_labels
    _UpperCamelCase:  List[Any]    						=			vocab_size
    _UpperCamelCase:  Dict    						=			hidden_size
    _UpperCamelCase:  int    						=			num_hidden_layers
    _UpperCamelCase:  Dict    						=			num_attention_heads
    _UpperCamelCase:  Dict    						=			intermediate_size

    _UpperCamelCase:  Union[str, Any]    						=			hidden_dropout_prob
    _UpperCamelCase:  Any    						=			attention_probs_dropout_prob
    _UpperCamelCase:  Optional[Any]    						=			max_position_embeddings
    _UpperCamelCase:  Union[str, Any]    						=			eos_token_id
    _UpperCamelCase:  Union[str, Any]    						=			pad_token_id
    _UpperCamelCase:  Optional[Any]    						=			bos_token_id
    _UpperCamelCase:  Optional[Any]    						=			attention_window

    # `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
    # [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
    # returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
    # because its local attention only attends to `self.attention_window` and one before and one after
    _UpperCamelCase:  Any    						=			self.attention_window + 2

    # because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
    # the `test_attention_outputs` and `test_hidden_states_output` tests
    _UpperCamelCase:  Tuple    						=			(
        self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
    )



   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  Union[str, Any]    						=			ids_tensor([self.batch_size, self.seq_length - 1]		,	self.vocab_size				)
    _UpperCamelCase:  Union[str, Any]    						=			tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size				)		,	1				)
    _UpperCamelCase:  Union[str, Any]    						=			tf.concat([input_ids, eos_tensor]		,	axis=1				)

    _UpperCamelCase:  Union[str, Any]    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size				)

    _UpperCamelCase:  Any    						=			self.config_cls(
        vocab_size=self.vocab_size		,	d_model=self.hidden_size		,	encoder_layers=self.num_hidden_layers		,	decoder_layers=self.num_hidden_layers		,	encoder_attention_heads=self.num_attention_heads		,	decoder_attention_heads=self.num_attention_heads		,	encoder_ffn_dim=self.intermediate_size		,	decoder_ffn_dim=self.intermediate_size		,	dropout=self.hidden_dropout_prob		,	attention_dropout=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	eos_token_ids=[2]		,	bos_token_id=self.bos_token_id		,	pad_token_id=self.pad_token_id		,	decoder_start_token_id=self.pad_token_id		,	attention_window=self.attention_window		,	**self.config_updates		,	)
    _UpperCamelCase:  Dict    						=			prepare_led_inputs_dict(_lowercase		,	_lowercase		,	_lowercase				)
    _UpperCamelCase:  Tuple    						=			tf.concat(
        [tf.zeros_like(_lowercase				)[:, :-1], tf.ones_like(_lowercase				)[:, -1:]]		,	axis=-1		,	)
    _UpperCamelCase:  List[Any]    						=			global_attention_mask
    return config, inputs_dict



   def 	lowerCAmelCase    (						self       :	Dict		,	_lowercase       :	List[str]		,	_lowercase       :	Union[str, Any]				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			TFLEDModel(config=_lowercase				).get_decoder()
    _UpperCamelCase:  Optional[int]    						=			inputs_dict['''input_ids''']

    _UpperCamelCase:  Dict    						=			input_ids[:1, :]
    _UpperCamelCase:  Optional[Any]    						=			inputs_dict['''attention_mask'''][:1, :]
    _UpperCamelCase:  Optional[int]    						=			1

    # first forward pass
    _UpperCamelCase:  Optional[Any]    						=			model(_lowercase		,	attention_mask=_lowercase		,	use_cache=_lowercase				)

    _UpperCamelCase						,				_UpperCamelCase:  Union[str, Any]    						=			outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    _UpperCamelCase:  Any    						=			ids_tensor((self.batch_size, 3)		,	config.vocab_size				)
    _UpperCamelCase:  List[str]    						=			tf.cast(ids_tensor((self.batch_size, 3)		,	2				)		,	tf.inta				)

    # append to next input_ids and
    _UpperCamelCase:  Optional[int]    						=			tf.concat([input_ids, next_tokens]		,	axis=-1				)
    _UpperCamelCase:  Union[str, Any]    						=			tf.concat([attention_mask, next_attn_mask]		,	axis=-1				)

    _UpperCamelCase:  str    						=			model(_lowercase		,	attention_mask=_lowercase				)[0]
    _UpperCamelCase:  Optional[int]    						=			model(_lowercase		,	attention_mask=_lowercase		,	past_key_values=_lowercase				)[0]

    self.parent.assertEqual(next_tokens.shape[1]		,	output_from_past.shape[1]				)

    # select random slice
    _UpperCamelCase:  Dict    						=			int(ids_tensor((1,)		,	output_from_past.shape[-1]				)				)
    _UpperCamelCase:  Any    						=			output_from_no_past[:, -3:, random_slice_idx]
    _UpperCamelCase:  Optional[Any]    						=			output_from_past[:, :, random_slice_idx]

    # test that outputs are equal for slice
    tf.debugging.assert_near(_lowercase		,	_lowercase		,	rtol=1E-3				)




def    lowerCAmelCase_  (		lowercase:				Tuple	,       lowercase:				List[str]	,       lowercase:				List[str]	,       lowercase:				int=None	,       lowercase:				int=None	,       lowercase:				List[str]=None	,       lowercase:				List[str]=None	,       )     ->					List[str]:







 '''simple docstring'''
 if attention_mask is None:
  _UpperCamelCase:  Optional[Any]    						=			tf.cast(tf.math.not_equal(lowercase	,       config.pad_token_id		)	,       tf.inta		)
 if decoder_attention_mask is None:
  _UpperCamelCase:  int    						=			tf.concat(
      [
          tf.ones(decoder_input_ids[:, :1].shape	,       dtype=tf.inta		),
          tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]	,       config.pad_token_id		)	,       tf.inta		),
      ]	,       axis=-1	,       )
 if head_mask is None:
  _UpperCamelCase:  str    						=			tf.ones((config.encoder_layers, config.encoder_attention_heads)		)
 if decoder_head_mask is None:
  _UpperCamelCase:  Optional[Any]    						=			tf.ones((config.decoder_layers, config.decoder_attention_heads)		)
 return {
     "input_ids": input_ids,
     "attention_mask": attention_mask,
     "decoder_input_ids": decoder_input_ids,
     "decoder_attention_mask": decoder_attention_mask,
     "head_mask": head_mask,
     "decoder_head_mask": decoder_head_mask,
 }




@require_tf
class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Optional[int]	      = (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
   lowerCAmelCase   :     List[str]	      = (TFLEDForConditionalGeneration,) if is_tf_available() else ()
   lowerCAmelCase   :     Union[str, Any]	      = (
       {
           '''conversational''': TFLEDForConditionalGeneration,
           '''feature-extraction''': TFLEDModel,
           '''summarization''': TFLEDForConditionalGeneration,
           '''text2text-generation''': TFLEDForConditionalGeneration,
           '''translation''': TFLEDForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   lowerCAmelCase   :     str	      = True
   lowerCAmelCase   :     Any	      = False
   lowerCAmelCase   :     Union[str, Any]	      = False
   lowerCAmelCase   :     Any	      = False



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			TFLEDModelTester(self				)
    _UpperCamelCase:  Any    						=			ConfigTester(self		,	config_class=_lowercase				)



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    self.config_tester.run_common_tests()



   def 	lowerCAmelCase    (						self       :	Tuple				):



    """simple docstring"""



    _UpperCamelCase:  Union[str, Any]    						=			self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*_lowercase				)



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase						,				_UpperCamelCase:  str    						=			self.model_tester.prepare_config_and_inputs_for_common()
    _UpperCamelCase:  str    						=			tf.zeros_like(inputs_dict['''attention_mask''']				)
    _UpperCamelCase:  Any    						=			2
    _UpperCamelCase:  List[str]    						=			tf.where(
        tf.range(self.model_tester.seq_length				)[None, :] < num_global_attn_indices		,	1		,	inputs_dict['''global_attention_mask''']		,	)

    _UpperCamelCase:  int    						=			True
    _UpperCamelCase:  str    						=			self.model_tester.seq_length
    _UpperCamelCase:  Dict    						=			self.model_tester.encoder_seq_length

    def check_decoder_attentions_output(_lowercase       :	Union[str, Any]				):
     _UpperCamelCase:  Tuple    						=			outputs.decoder_attentions
     self.assertEqual(len(_lowercase				)		,	self.model_tester.num_hidden_layers				)
     self.assertListEqual(
         list(decoder_attentions[0].shape[-3:]				)		,	[self.model_tester.num_attention_heads, seq_length, seq_length]		,	)

    def check_encoder_attentions_output(_lowercase       :	int				):
     _UpperCamelCase:  str    						=			[t.numpy() for t in outputs.encoder_attentions]
     _UpperCamelCase:  Any    						=			[t.numpy() for t in outputs.encoder_global_attentions]
     self.assertEqual(len(_lowercase				)		,	self.model_tester.num_hidden_layers				)
     self.assertEqual(len(_lowercase				)		,	self.model_tester.num_hidden_layers				)
     self.assertListEqual(
         list(attentions[0].shape[-3:]				)		,	[self.model_tester.num_attention_heads, seq_length, seq_length]		,	)
     self.assertListEqual(
         list(global_attentions[0].shape[-3:]				)		,	[self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]		,	)

    for model_class in self.all_model_classes:
     _UpperCamelCase:  Optional[int]    						=			True
     _UpperCamelCase:  List[Any]    						=			False
     _UpperCamelCase:  Optional[int]    						=			False
     _UpperCamelCase:  Union[str, Any]    						=			model_class(_lowercase				)
     _UpperCamelCase:  Optional[int]    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)
     _UpperCamelCase:  List[str]    						=			len(_lowercase				)
     self.assertEqual(config.output_hidden_states		,	_lowercase				)
     check_encoder_attentions_output(_lowercase				)

     if self.is_encoder_decoder:
      _UpperCamelCase:  Optional[int]    						=			model_class(_lowercase				)
      _UpperCamelCase:  Union[str, Any]    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)
      self.assertEqual(config.output_hidden_states		,	_lowercase				)
      check_decoder_attentions_output(_lowercase				)

     # Check that output attentions can also be changed via the config
     del inputs_dict["output_attentions"]
     _UpperCamelCase:  Union[str, Any]    						=			True
     _UpperCamelCase:  List[str]    						=			model_class(_lowercase				)
     _UpperCamelCase:  Dict    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)
     self.assertEqual(config.output_hidden_states		,	_lowercase				)
     check_encoder_attentions_output(_lowercase				)

     # Check attention is always last and order is fine
     _UpperCamelCase:  str    						=			True
     _UpperCamelCase:  Dict    						=			True
     _UpperCamelCase:  int    						=			model_class(_lowercase				)
     _UpperCamelCase:  Optional[int]    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)

     self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)		,	len(_lowercase				)				)
     self.assertEqual(model.config.output_hidden_states		,	_lowercase				)
     check_encoder_attentions_output(_lowercase				)



   @unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''				)
   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    pass



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    pass




def    lowerCAmelCase_  (		lowercase:				Tuple		)     ->					List[str]:







 '''simple docstring'''
 return tf.constant(lowercase	,       dtype=tf.intaa		)


UpperCAmelCase_           =				1E-4





@slow
@require_tf
class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   def 	lowerCAmelCase    (						self       :	str				):



    """simple docstring"""



    _UpperCamelCase:  List[str]    						=			TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''				).led

    # change to intended input here
    _UpperCamelCase:  List[str]    						=			_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  int    						=			_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  Union[str, Any]    						=			prepare_led_inputs_dict(model.config		,	_lowercase		,	_lowercase				)
    _UpperCamelCase:  Optional[Any]    						=			model(**_lowercase				)[0]
    _UpperCamelCase:  Optional[int]    						=			(1, 1_024, 768)
    self.assertEqual(output.shape		,	_lowercase				)
    # change to expected output here
    _UpperCamelCase:  Optional[Any]    						=			tf.convert_to_tensor(
        [[2.3050, 2.8279, 0.6531], [-1.8457, -0.1455, -3.5661], [-1.0186, 0.4586, -2.2043]]		,	)
    tf.debugging.assert_near(output[:, :3, :3]		,	_lowercase		,	atol=1E-3				)



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



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''				)

    # change to intended input here
    _UpperCamelCase:  Dict    						=			_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  List[str]    						=			_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  Tuple    						=			prepare_led_inputs_dict(model.config		,	_lowercase		,	_lowercase				)
    _UpperCamelCase:  Any    						=			model(**_lowercase				)[0]
    _UpperCamelCase:  Dict    						=			(1, 1_024, model.config.vocab_size)
    self.assertEqual(output.shape		,	_lowercase				)
    # change to expected output here
    _UpperCamelCase:  Optional[Any]    						=			tf.convert_to_tensor(
        [[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]]		,	)
    tf.debugging.assert_near(output[:, :3, :3]		,	_lowercase		,	atol=1E-3		,	rtol=1E-3				)