"""simple docstring"""






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  _A						:
      snake_case__			:   List[Any]							  =	LEDConfig
      snake_case__			:   List[str]							  =	{}
      snake_case__			:   int							  =	'gelu'

      def __init__(						self		,       __lowerCAmelCase		,       __lowerCAmelCase=13		,       __lowerCAmelCase=7		,       __lowerCAmelCase=True		,       __lowerCAmelCase=False		,       __lowerCAmelCase=99		,       __lowerCAmelCase=32		,       __lowerCAmelCase=2		,       __lowerCAmelCase=4		,       __lowerCAmelCase=37		,       __lowerCAmelCase=0.1		,       __lowerCAmelCase=0.1		,       __lowerCAmelCase=20		,       __lowerCAmelCase=2		,       __lowerCAmelCase=1		,       __lowerCAmelCase=0		,       __lowerCAmelCase=4		,       ):







            """simple docstring"""



            lowercase														=	parent
            lowercase														=	batch_size
            lowercase														=	seq_length
            lowercase														=	is_training
            lowercase														=	use_labels
            lowercase														=	vocab_size
            lowercase														=	hidden_size
            lowercase														=	num_hidden_layers
            lowercase														=	num_attention_heads
            lowercase														=	intermediate_size

            lowercase														=	hidden_dropout_prob
            lowercase														=	attention_probs_dropout_prob
            lowercase														=	max_position_embeddings
            lowercase														=	eos_token_id
            lowercase														=	pad_token_id
            lowercase														=	bos_token_id
            lowercase														=	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
            lowercase														=	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
            lowercase														=	(
                self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
            )

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	ids_tensor([self.batch_size, self.seq_length - 1]		,       self.vocab_size						)
            lowercase														=	tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size						)		,       1						)
            lowercase														=	tf.concat([input_ids, eos_tensor]		,       axis=1						)

            lowercase														=	ids_tensor([self.batch_size, self.seq_length]		,       self.vocab_size						)

            lowercase														=	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		,       )
            lowercase														=	prepare_led_inputs_dict(__lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase						)
            lowercase														=	tf.concat(
                [tf.zeros_like(__lowerCAmelCase						)[:, :-1], tf.ones_like(__lowerCAmelCase						)[:, -1:]]		,       axis=-1		,       )
            lowercase														=	global_attention_mask
            return config, inputs_dict







      def 				A__       (						self		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



            lowercase														=	TFLEDModel(config=__lowerCAmelCase						).get_decoder()
            lowercase														=	inputs_dict["""input_ids"""]

            lowercase														=	input_ids[:1, :]
            lowercase														=	inputs_dict["""attention_mask"""][:1, :]
            lowercase														=	1

            # first forward pass
            lowercase														=	model(__lowerCAmelCase		,       attention_mask=__lowerCAmelCase		,       use_cache=__lowerCAmelCase						)

            lowercase							,							lowercase														=	outputs.to_tuple()

            # create hypothetical next token and extent to next_input_ids
            lowercase														=	ids_tensor((self.batch_size, 3)		,       config.vocab_size						)
            lowercase														=	tf.cast(ids_tensor((self.batch_size, 3)		,       2						)		,       tf.inta						)

            # append to next input_ids and
            lowercase														=	tf.concat([input_ids, next_tokens]		,       axis=-1						)
            lowercase														=	tf.concat([attention_mask, next_attn_mask]		,       axis=-1						)

            lowercase														=	model(__lowerCAmelCase		,       attention_mask=__lowerCAmelCase						)[0]
            lowercase														=	model(__lowerCAmelCase		,       attention_mask=__lowerCAmelCase		,       past_key_values=__lowerCAmelCase						)[0]

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

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

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






def 			UpperCAmelCase__       (				lowerCAmelCase__     :Dict					,		lowerCAmelCase__     :Tuple					,		lowerCAmelCase__     :str					,		lowerCAmelCase__     :Any=None					,		lowerCAmelCase__     :Tuple=None					,		lowerCAmelCase__     :Optional[Any]=None					,		lowerCAmelCase__     :str=None					,		) ->    List[str]:
      '''simple docstring'''


      if attention_mask is None:
            lowercase														=	tf.cast(tf.math.not_equal(lowerCAmelCase__					,		config.pad_token_id		)					,		tf.inta		)
      if decoder_attention_mask is None:
            lowercase														=	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:
            lowercase														=	tf.ones((config.encoder_layers, config.encoder_attention_heads)		)
      if decoder_head_mask is None:
            lowercase														=	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  _A						(       lowerCAmelCase  ,    lowerCAmelCase  ,    unittest.TestCase	):
      snake_case__			:   Union[str, Any]							  =	(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
      snake_case__			:   Any							  =	(TFLEDForConditionalGeneration,) if is_tf_available() else ()
      snake_case__			:   Optional[int]							  =	(
          {
              'conversational': TFLEDForConditionalGeneration,
              'feature-extraction': TFLEDModel,
              'summarization': TFLEDForConditionalGeneration,
              'text2text-generation': TFLEDForConditionalGeneration,
              'translation': TFLEDForConditionalGeneration,
          }
          if is_tf_available()
          else {}
      )
      snake_case__			:   Dict							  =	True
      snake_case__			:   Optional[int]							  =	False
      snake_case__			:   List[str]							  =	False
      snake_case__			:   Union[str, Any]							  =	False

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	TFLEDModelTester(self						)
            lowercase														=	ConfigTester(self		,       config_class=__lowerCAmelCase						)

      def 				A__       (						self						):







            """simple docstring"""



            self.config_tester.run_common_tests()

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.check_decoder_model_past_large_inputs(*__lowerCAmelCase						)

      def 				A__       (						self						):







            """simple docstring"""



            lowercase							,							lowercase														=	self.model_tester.prepare_config_and_inputs_for_common()
            lowercase														=	tf.zeros_like(inputs_dict["""attention_mask"""]						)
            lowercase														=	2
            lowercase														=	tf.where(
                tf.range(self.model_tester.seq_length						)[None, :] < num_global_attn_indices		,       1		,       inputs_dict["""global_attention_mask"""]		,       )

            lowercase														=	True
            lowercase														=	self.model_tester.seq_length
            lowercase														=	self.model_tester.encoder_seq_length

            def check_decoder_attentions_output(__lowerCAmelCase						):
                  lowercase														=	outputs.decoder_attentions
                  self.assertEqual(len(__lowerCAmelCase						)		,       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(__lowerCAmelCase						):
                  lowercase														=	[t.numpy() for t in outputs.encoder_attentions]
                  lowercase														=	[t.numpy() for t in outputs.encoder_global_attentions]
                  self.assertEqual(len(__lowerCAmelCase						)		,       self.model_tester.num_hidden_layers						)
                  self.assertEqual(len(__lowerCAmelCase						)		,       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:
                  lowercase														=	True
                  lowercase														=	False
                  lowercase														=	False
                  lowercase														=	model_class(__lowerCAmelCase						)
                  lowercase														=	model(self._prepare_for_class(__lowerCAmelCase		,       __lowerCAmelCase						)						)
                  lowercase														=	len(__lowerCAmelCase						)
                  self.assertEqual(config.output_hidden_states		,       __lowerCAmelCase						)
                  check_encoder_attentions_output(__lowerCAmelCase						)

                  if self.is_encoder_decoder:
                        lowercase														=	model_class(__lowerCAmelCase						)
                        lowercase														=	model(self._prepare_for_class(__lowerCAmelCase		,       __lowerCAmelCase						)						)
                        self.assertEqual(config.output_hidden_states		,       __lowerCAmelCase						)
                        check_decoder_attentions_output(__lowerCAmelCase						)

                  # Check that output attentions can also be changed via the config
                  del inputs_dict["output_attentions"]
                  lowercase														=	True
                  lowercase														=	model_class(__lowerCAmelCase						)
                  lowercase														=	model(self._prepare_for_class(__lowerCAmelCase		,       __lowerCAmelCase						)						)
                  self.assertEqual(config.output_hidden_states		,       __lowerCAmelCase						)
                  check_encoder_attentions_output(__lowerCAmelCase						)

                  # Check attention is always last and order is fine
                  lowercase														=	True
                  lowercase														=	True
                  lowercase														=	model_class(__lowerCAmelCase						)
                  lowercase														=	model(self._prepare_for_class(__lowerCAmelCase		,       __lowerCAmelCase						)						)

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

      @unittest.skip("""LED keeps using potentially symbolic tensors in conditionals and breaks tracing."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass







      def 				A__       (						self						):







            """simple docstring"""



            pass






def 			UpperCAmelCase__       (				lowerCAmelCase__     :List[str]		) ->    Tuple:
      '''simple docstring'''


      return tf.constant(lowerCAmelCase__					,		dtype=tf.intaa		)


__lowerCAmelCase    :       Dict					     =1E-4






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

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""						).led

            # change to intended input here
            lowercase														=	_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]						)
            lowercase														=	_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]						)
            lowercase														=	prepare_led_inputs_dict(model.config		,       __lowerCAmelCase		,       __lowerCAmelCase						)
            lowercase														=	model(**__lowerCAmelCase						)[0]
            lowercase														=	(1, 1024, 768)
            self.assertEqual(output.shape		,       __lowerCAmelCase						)
            # change to expected output here
            lowercase														=	tf.convert_to_tensor(
                [[2.3_0_5_0, 2.8_2_7_9, 0.6_5_3_1], [-1.8_4_5_7, -0.1_4_5_5, -3.5_6_6_1], [-1.0_1_8_6, 0.4_5_8_6, -2.2_0_4_3]]		,       )
            tf.debugging.assert_near(output[:, :3, :3]		,       __lowerCAmelCase		,       atol=1E-3						)







      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""						)

            # change to intended input here
            lowercase														=	_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]						)
            lowercase														=	_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]						)
            lowercase														=	prepare_led_inputs_dict(model.config		,       __lowerCAmelCase		,       __lowerCAmelCase						)
            lowercase														=	model(**__lowerCAmelCase						)[0]
            lowercase														=	(1, 1024, model.config.vocab_size)
            self.assertEqual(output.shape		,       __lowerCAmelCase						)
            # change to expected output here
            lowercase														=	tf.convert_to_tensor(
                [[3_3.6_5_0_7, 6.4_5_7_2, 1_6.8_0_8_9], [5.8_7_3_9, -2.4_2_3_8, 1_1.2_9_0_2], [-3.2_1_3_9, -4.3_1_4_9, 4.2_7_8_3]]		,       )
            tf.debugging.assert_near(output[:, :3, :3]		,       __lowerCAmelCase		,       atol=1E-3		,       rtol=1E-3						)




"""simple docstring"""






import logging
import random

import ray

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


__lowerCAmelCase    :       int					     =logging.getLogger(__name__)





class  _A						:

      def __init__(						self						):







            """simple docstring"""



            lowercase														=	False

      def 				A__       (						self		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



            if not self.initialized:
                  lowercase														=	RagRetriever(
                      __lowerCAmelCase		,       question_encoder_tokenizer=__lowerCAmelCase		,       generator_tokenizer=__lowerCAmelCase		,       index=__lowerCAmelCase		,       init_retrieval=__lowerCAmelCase		,       )
                  lowercase														=	True

      def 				A__       (						self						):







            """simple docstring"""



            self.retriever.index.init_index()







      def 				A__       (						self		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



            lowercase							,							lowercase														=	self.retriever._main_retrieve(__lowerCAmelCase		,       __lowerCAmelCase						)
            return doc_ids, retrieved_doc_embeds






class  _A						(       lowerCAmelCase	):

      def __init__(						self		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase=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		,       )
            lowercase														=	retrieval_workers
            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						):







            """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		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



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

      @classmethod
      def 				A__       (						cls		,       __lowerCAmelCase		,       __lowerCAmelCase=None		,       **__lowerCAmelCase						):







            """simple docstring"""



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







      @classmethod
      def 				A__       (						cls		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase=None		,       **__lowerCAmelCase						):







            """simple docstring"""



            lowercase														=	kwargs.pop("""config"""		,       __lowerCAmelCase						) or RagConfig.from_pretrained(__lowerCAmelCase		,       **__lowerCAmelCase						)
            lowercase														=	RagTokenizer.from_pretrained(__lowerCAmelCase		,       config=__lowerCAmelCase						)
            lowercase														=	rag_tokenizer.question_encoder
            lowercase														=	rag_tokenizer.generator
            if indexed_dataset is not None:
                  lowercase														=	"""custom"""
                  lowercase														=	CustomHFIndex(config.retrieval_vector_size		,       __lowerCAmelCase						)
            else:
                  lowercase														=	cls._build_index(__lowerCAmelCase						)
            return cls(
                __lowerCAmelCase		,       question_encoder_tokenizer=__lowerCAmelCase		,       generator_tokenizer=__lowerCAmelCase		,       retrieval_workers=__lowerCAmelCase		,       index=__lowerCAmelCase		,       )



