





'''simple docstring'''

import argparse
import csv
import logging
import os
import random

import numpy as np
import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from tqdm import tqdm, trange

from transformers import (
    CONFIG_NAME,
    WEIGHTS_NAME,
    AdamW,
    OpenAIGPTDoubleHeadsModel,
    OpenAIGPTTokenizer,
    get_linear_schedule_with_warmup,
)


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
UpperCAmelCase					:		List[str]		=							logging.getLogger(__name__)




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_	):





      """simple docstring"""



      _snake_case		:	Optional[int]					     =     np.argmax(lowerCAmelCase_			,      axis=1	)
      return np.sum(outputs == labels	)




def    _a    (							lowerCAmelCase_	):





      """simple docstring"""



      with open(lowerCAmelCase_			,      encoding='''utf_8'''	) as f:
            _snake_case		:	Dict					     =     csv.reader(lowerCAmelCase_	)
            _snake_case		:	Tuple					     =     []
            next(lowerCAmelCase_	)  # skip the first line
            for line in tqdm(lowerCAmelCase_	):
                  output.append((''' '''.join(line[1:5]	), line[5], line[6], int(line[-1]	) - 1)	)
      return output




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_	):





      """simple docstring"""



      _snake_case		:	Union[str, Any]					     =     []
      for dataset in encoded_datasets:
            _snake_case		:	List[Any]					     =     len(lowerCAmelCase_	)
            _snake_case		:	int					     =     np.zeros((n_batch, 2, input_len)			,      dtype=np.intaa	)
            _snake_case		:	Optional[Any]					     =     np.zeros((n_batch, 2)			,      dtype=np.intaa	)
            _snake_case		:	Optional[Any]					     =     np.full((n_batch, 2, input_len)			,      fill_value=-100			,      dtype=np.intaa	)
            _snake_case		:	int					     =     np.zeros((n_batch,)			,      dtype=np.intaa	)
            for (
                i,
                (story, conta, conta, mc_label),
            ) in enumerate(lowerCAmelCase_	):
                  _snake_case		:	Any					     =     [start_token] + story[:cap_length] + [delimiter_token] + conta[:cap_length] + [clf_token]
                  _snake_case		:	List[Any]					     =     [start_token] + story[:cap_length] + [delimiter_token] + conta[:cap_length] + [clf_token]
                  _snake_case		:	List[Any]					     =     with_conta
                  _snake_case		:	Dict					     =     with_conta
                  _snake_case		:	str					     =     len(lowerCAmelCase_	) - 1
                  _snake_case		:	List[Any]					     =     len(lowerCAmelCase_	) - 1
                  _snake_case		:	Optional[Any]					     =     with_conta
                  _snake_case		:	List[Any]					     =     with_conta
                  _snake_case		:	Optional[Any]					     =     mc_label
            _snake_case		:	Optional[Any]					     =     (input_ids, mc_token_ids, lm_labels, mc_labels)
            tensor_datasets.append(tuple(torch.tensor(lowerCAmelCase_	) for t in all_inputs	)	)
      return tensor_datasets




def    _a    (							):





      """simple docstring"""



      _snake_case		:	List[Any]					     =     argparse.ArgumentParser()
      parser.add_argument('''--model_name'''			,      type=lowerCAmelCase_			,      default='''openai-gpt'''			,      help='''pretrained model name'''	)
      parser.add_argument('''--do_train'''			,      action='''store_true'''			,      help='''Whether to run training.'''	)
      parser.add_argument('''--do_eval'''			,      action='''store_true'''			,      help='''Whether to run eval on the dev set.'''	)
      parser.add_argument(
          '''--output_dir'''			,      default=lowerCAmelCase_			,      type=lowerCAmelCase_			,      required=lowerCAmelCase_			,      help='''The output directory where the model predictions and checkpoints will be written.'''			,      )
      parser.add_argument('''--train_dataset'''			,      type=lowerCAmelCase_			,      default=''''''	)
      parser.add_argument('''--eval_dataset'''			,      type=lowerCAmelCase_			,      default=''''''	)
      parser.add_argument('''--seed'''			,      type=lowerCAmelCase_			,      default=42	)
      parser.add_argument('''--num_train_epochs'''			,      type=lowerCAmelCase_			,      default=3	)
      parser.add_argument('''--train_batch_size'''			,      type=lowerCAmelCase_			,      default=8	)
      parser.add_argument('''--eval_batch_size'''			,      type=lowerCAmelCase_			,      default=16	)
      parser.add_argument('''--adam_epsilon'''			,      default=1E-8			,      type=lowerCAmelCase_			,      help='''Epsilon for Adam optimizer.'''	)
      parser.add_argument('''--max_grad_norm'''			,      type=lowerCAmelCase_			,      default=1	)
      parser.add_argument(
          '''--max_steps'''			,      default=-1			,      type=lowerCAmelCase_			,      help=(
              '''If > 0: set total number of training                         steps to perform. Override num_train_epochs.'''
          )			,      )
      parser.add_argument(
          '''--gradient_accumulation_steps'''			,      type=lowerCAmelCase_			,      default=1			,      help='''Number of updates steps to accumulate before                        performing a backward/update pass.'''			,      )
      parser.add_argument('''--learning_rate'''			,      type=lowerCAmelCase_			,      default=6.25E-5	)
      parser.add_argument('''--warmup_steps'''			,      default=0			,      type=lowerCAmelCase_			,      help='''Linear warmup over warmup_steps.'''	)
      parser.add_argument('''--lr_schedule'''			,      type=lowerCAmelCase_			,      default='''warmup_linear'''	)
      parser.add_argument('''--weight_decay'''			,      type=lowerCAmelCase_			,      default=0.01	)
      parser.add_argument('''--lm_coef'''			,      type=lowerCAmelCase_			,      default=0.9	)
      parser.add_argument('''--n_valid'''			,      type=lowerCAmelCase_			,      default=374	)

      parser.add_argument('''--server_ip'''			,      type=lowerCAmelCase_			,      default=''''''			,      help='''Can be used for distant debugging.'''	)
      parser.add_argument('''--server_port'''			,      type=lowerCAmelCase_			,      default=''''''			,      help='''Can be used for distant debugging.'''	)
      _snake_case		:	Optional[int]					     =     parser.parse_args()
      print(lowerCAmelCase_	)

      if args.server_ip and args.server_port:
            # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
            import ptvsd

            print('''Waiting for debugger attach'''	)
            ptvsd.enable_attach(address=(args.server_ip, args.server_port)			,      redirect_output=lowerCAmelCase_	)
            ptvsd.wait_for_attach()

      random.seed(args.seed	)
      np.random.seed(args.seed	)
      torch.manual_seed(args.seed	)
      torch.cuda.manual_seed_all(args.seed	)

      _snake_case		:	Optional[Any]					     =     torch.device('''cuda''' if torch.cuda.is_available() else '''cpu'''	)
      _snake_case		:	Dict					     =     torch.cuda.device_count()
      logger.info('''device: {}, n_gpu {}'''.format(lowerCAmelCase_			,      lowerCAmelCase_	)	)

      if not args.do_train and not args.do_eval:
            raise ValueError('''At least one of `do_train` or `do_eval` must be True.'''	)

      if not os.path.exists(args.output_dir	):
            os.makedirs(args.output_dir	)

      # Load tokenizer and model
      # This loading functions also add new tokens and embeddings called `special tokens`
      # These new embeddings will be fine-tuned on the RocStories dataset
      _snake_case		:	Optional[Any]					     =     ['''_start_''', '''_delimiter_''', '''_classify_''']
      _snake_case		:	Optional[int]					     =     OpenAIGPTTokenizer.from_pretrained(args.model_name	)
      tokenizer.add_tokens(lowerCAmelCase_	)
      _snake_case		:	Optional[Any]					     =     tokenizer.convert_tokens_to_ids(lowerCAmelCase_	)
      _snake_case		:	List[str]					     =     OpenAIGPTDoubleHeadsModel.from_pretrained(args.model_name	)
      model.resize_token_embeddings(len(lowerCAmelCase_	)	)
      model.to(lowerCAmelCase_	)

      # Load and encode the datasets
      def tokenize_and_encode(lowerCAmelCase_	):
            if isinstance(lowerCAmelCase_			,      lowerCAmelCase_	):
                  return tokenizer.convert_tokens_to_ids(tokenizer.tokenize(lowerCAmelCase_	)	)
            elif isinstance(lowerCAmelCase_			,      lowerCAmelCase_	):
                  return obj
            return [tokenize_and_encode(lowerCAmelCase_	) for o in obj]

      logger.info('''Encoding dataset...'''	)
      _snake_case		:	Optional[Any]					     =     load_rocstories_dataset(args.train_dataset	)
      _snake_case		:	List[Any]					     =     load_rocstories_dataset(args.eval_dataset	)
      _snake_case		:	Dict					     =     (train_dataset, eval_dataset)
      _snake_case		:	Optional[Any]					     =     tokenize_and_encode(lowerCAmelCase_	)

      # Compute the max input length for the Transformer
      _snake_case		:	Tuple					     =     model.config.n_positions // 2 - 2
      _snake_case		:	Optional[Any]					     =     max(
          len(story[:max_length]	) + max(len(conta[:max_length]	)			,      len(conta[:max_length]	)	) + 3
          for dataset in encoded_datasets
          for story, conta, conta, _ in dataset	)
      _snake_case		:	Optional[int]					     =     min(lowerCAmelCase_			,      model.config.n_positions	)  # Max size of input for the pre-trained model

      # Prepare inputs tensors and dataloaders
      _snake_case		:	int					     =     pre_process_datasets(lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_			,      *lowerCAmelCase_	)
      _snake_case     ,							_snake_case		:	Dict					     =     tensor_datasets[0], tensor_datasets[1]

      _snake_case		:	str					     =     TensorDataset(*lowerCAmelCase_	)
      _snake_case		:	List[Any]					     =     RandomSampler(lowerCAmelCase_	)
      _snake_case		:	List[Any]					     =     DataLoader(lowerCAmelCase_			,      sampler=lowerCAmelCase_			,      batch_size=args.train_batch_size	)

      _snake_case		:	str					     =     TensorDataset(*lowerCAmelCase_	)
      _snake_case		:	Any					     =     SequentialSampler(lowerCAmelCase_	)
      _snake_case		:	str					     =     DataLoader(lowerCAmelCase_			,      sampler=lowerCAmelCase_			,      batch_size=args.eval_batch_size	)

      # Prepare optimizer
      if args.do_train:
            if args.max_steps > 0:
                  _snake_case		:	Optional[int]					     =     args.max_steps
                  _snake_case		:	Optional[int]					     =     args.max_steps // (len(lowerCAmelCase_	) // args.gradient_accumulation_steps) + 1
            else:
                  _snake_case		:	Dict					     =     len(lowerCAmelCase_	) // args.gradient_accumulation_steps * args.num_train_epochs

            _snake_case		:	Union[str, Any]					     =     list(model.named_parameters()	)
            _snake_case		:	str					     =     ['''bias''', '''LayerNorm.bias''', '''LayerNorm.weight''']
            _snake_case		:	Union[str, Any]					     =     [
                {
                    '''params''': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay	)],
                    '''weight_decay''': args.weight_decay,
                },
                {'''params''': [p for n, p in param_optimizer if any(nd in n for nd in no_decay	)], '''weight_decay''': 0.0},
            ]
            _snake_case		:	Any					     =     AdamW(lowerCAmelCase_			,      lr=args.learning_rate			,      eps=args.adam_epsilon	)
            _snake_case		:	str					     =     get_linear_schedule_with_warmup(
                lowerCAmelCase_			,      num_warmup_steps=args.warmup_steps			,      num_training_steps=lowerCAmelCase_	)

      if args.do_train:
            _snake_case     ,							_snake_case     ,							_snake_case		:	Optional[int]					     =     0, 0, None
            model.train()
            for _ in trange(int(args.num_train_epochs	)			,      desc='''Epoch'''	):
                  _snake_case		:	str					     =     0
                  _snake_case		:	List[str]					     =     0
                  _snake_case		:	Tuple					     =     tqdm(lowerCAmelCase_			,      desc='''Training'''	)
                  for step, batch in enumerate(lowerCAmelCase_	):
                        _snake_case		:	Any					     =     tuple(t.to(lowerCAmelCase_	) for t in batch	)
                        _snake_case     ,							_snake_case     ,							_snake_case     ,							_snake_case		:	Optional[Any]					     =     batch
                        _snake_case		:	Union[str, Any]					     =     model(lowerCAmelCase_			,      mc_token_ids=lowerCAmelCase_			,      lm_labels=lowerCAmelCase_			,      mc_labels=lowerCAmelCase_	)
                        _snake_case		:	List[Any]					     =     args.lm_coef * losses[0] + losses[1]
                        loss.backward()
                        optimizer.step()
                        scheduler.step()
                        optimizer.zero_grad()
                        tr_loss += loss.item()
                        _snake_case		:	Optional[Any]					     =     (
                            loss.item() if exp_average_loss is None else 0.7 * exp_average_loss + 0.3 * loss.item()
                        )
                        nb_tr_steps += 1
                        _snake_case		:	int					     =     '''Training loss: {:.2e} lr: {:.2e}'''.format(lowerCAmelCase_			,      scheduler.get_lr()[0]	)

    # Save a trained model
      if args.do_train:
            # Save a trained model, configuration and tokenizer
            _snake_case		:	Optional[int]					     =     model.module if hasattr(lowerCAmelCase_			,      '''module'''	) else model  # Only save the model itself

            # If we save using the predefined names, we can load using `from_pretrained`
            _snake_case		:	Dict					     =     os.path.join(args.output_dir			,      lowerCAmelCase_	)
            _snake_case		:	Dict					     =     os.path.join(args.output_dir			,      lowerCAmelCase_	)

            torch.save(model_to_save.state_dict()			,      lowerCAmelCase_	)
            model_to_save.config.to_json_file(lowerCAmelCase_	)
            tokenizer.save_vocabulary(args.output_dir	)

            # Load a trained model and vocabulary that you have fine-tuned
            _snake_case		:	Dict					     =     OpenAIGPTDoubleHeadsModel.from_pretrained(args.output_dir	)
            _snake_case		:	Optional[int]					     =     OpenAIGPTTokenizer.from_pretrained(args.output_dir	)
            model.to(lowerCAmelCase_	)

      if args.do_eval:
            model.eval()
            _snake_case     ,							_snake_case		:	Tuple					     =     0, 0
            _snake_case     ,							_snake_case		:	str					     =     0, 0
            for batch in tqdm(lowerCAmelCase_			,      desc='''Evaluating'''	):
                  _snake_case		:	Any					     =     tuple(t.to(lowerCAmelCase_	) for t in batch	)
                  _snake_case     ,							_snake_case     ,							_snake_case     ,							_snake_case		:	Dict					     =     batch
                  with torch.no_grad():
                        _snake_case     ,							_snake_case     ,							_snake_case     ,							_snake_case		:	List[Any]					     =     model(
                            lowerCAmelCase_			,      mc_token_ids=lowerCAmelCase_			,      lm_labels=lowerCAmelCase_			,      mc_labels=lowerCAmelCase_	)

                  _snake_case		:	int					     =     mc_logits.detach().cpu().numpy()
                  _snake_case		:	Any					     =     mc_labels.to('''cpu'''	).numpy()
                  _snake_case		:	Optional[Any]					     =     accuracy(lowerCAmelCase_			,      lowerCAmelCase_	)

                  eval_loss += mc_loss.mean().item()
                  eval_accuracy += tmp_eval_accuracy

                  nb_eval_examples += input_ids.size(0	)
                  nb_eval_steps += 1

            _snake_case		:	Optional[int]					     =     eval_loss / nb_eval_steps
            _snake_case		:	Optional[int]					     =     eval_accuracy / nb_eval_examples
            _snake_case		:	List[Any]					     =     tr_loss / nb_tr_steps if args.do_train else None
            _snake_case		:	Tuple					     =     {'''eval_loss''': eval_loss, '''eval_accuracy''': eval_accuracy, '''train_loss''': train_loss}

            _snake_case		:	str					     =     os.path.join(args.output_dir			,      '''eval_results.txt'''	)
            with open(lowerCAmelCase_			,      '''w'''	) as writer:
                  logger.info('''***** Eval results *****'''	)
                  for key in sorted(result.keys()	):
                        logger.info('''  %s = %s'''			,      lowerCAmelCase_			,      str(result[key]	)	)
                        writer.write('''%s = %s\n''' % (key, str(result[key]	))	)


if __name__ == "__main__":
   main()









'''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 lowerCamelCase						:
  _lowercase  :	Any       					=				LEDConfig
  _lowercase  :	Any       					=				{}
  _lowercase  :	Optional[Any]       					=				"""gelu"""







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






        """simple docstring"""





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

        _snake_case		:	List[Any]					     =     hidden_dropout_prob
        _snake_case		:	List[str]					     =     attention_probs_dropout_prob
        _snake_case		:	Optional[int]					     =     max_position_embeddings
        _snake_case		:	Any					     =     eos_token_id
        _snake_case		:	List[Any]					     =     pad_token_id
        _snake_case		:	Optional[int]					     =     bos_token_id
        _snake_case		:	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
        _snake_case		:	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
        _snake_case		:	Tuple					     =     (
            self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
        )







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






        """simple docstring"""





        _snake_case		:	Optional[int]					     =     ids_tensor([self.batch_size, self.seq_length - 1]     ,     self.vocab_size		)
        _snake_case		:	Tuple					     =     tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)     ,     1		)
        _snake_case		:	Optional[int]					     =     tf.concat([input_ids, eos_tensor]     ,     axis=1		)

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

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







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__		)		->						int:






        """simple docstring"""





        _snake_case		:	int					     =     TFLEDModel(config=lowercase__		).get_decoder()
        _snake_case		:	Union[str, Any]					     =     inputs_dict['''input_ids''']

        _snake_case		:	List[str]					     =     input_ids[:1, :]
        _snake_case		:	Tuple					     =     inputs_dict['''attention_mask'''][:1, :]
        _snake_case		:	Dict					     =     1

        # first forward pass
        _snake_case		:	Optional[int]					     =     model(lowercase__     ,     attention_mask=lowercase__     ,     use_cache=lowercase__		)

        _snake_case     ,							_snake_case		:	Dict					     =     outputs.to_tuple()

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

        # append to next input_ids and
        _snake_case		:	Tuple					     =     tf.concat([input_ids, next_tokens]     ,     axis=-1		)
        _snake_case		:	List[Any]					     =     tf.concat([attention_mask, next_attn_mask]     ,     axis=-1		)

        _snake_case		:	List[Any]					     =     model(lowercase__     ,     attention_mask=lowercase__		)[0]
        _snake_case		:	Tuple					     =     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
        _snake_case		:	Tuple					     =     int(ids_tensor((1,)     ,     output_from_past.shape[-1]		)		)
        _snake_case		:	int					     =     output_from_no_past[:, -3:, random_slice_idx]
        _snake_case		:	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    _a    (							lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_=None			,      lowerCAmelCase_=None			,      lowerCAmelCase_=None			,      lowerCAmelCase_=None			,      ):





      """simple docstring"""



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







  def 				UpperCAmelCase_				(		self		)		->						Optional[Any]:






        """simple docstring"""





        _snake_case		:	str					     =     TFLEDModelTester(self		)
        _snake_case		:	Union[str, Any]					     =     ConfigTester(self     ,     config_class=lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						Tuple:






        """simple docstring"""





        self.config_tester.run_common_tests()







  def 				UpperCAmelCase_				(		self		)		->						List[str]:






        """simple docstring"""





        _snake_case		:	List[Any]					     =     self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_decoder_model_past_large_inputs(*lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





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

        _snake_case		:	Dict					     =     True
        _snake_case		:	str					     =     self.model_tester.seq_length
        _snake_case		:	Dict					     =     self.model_tester.encoder_seq_length

        def check_decoder_attentions_output(lowercase__		):
              _snake_case		:	Optional[int]					     =     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__		):
              _snake_case		:	int					     =     [t.numpy() for t in outputs.encoder_attentions]
              _snake_case		:	Tuple					     =     [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:
              _snake_case		:	Union[str, Any]					     =     True
              _snake_case		:	Dict					     =     False
              _snake_case		:	Union[str, Any]					     =     False
              _snake_case		:	List[Any]					     =     model_class(lowercase__		)
              _snake_case		:	Optional[Any]					     =     model(self._prepare_for_class(lowercase__     ,     lowercase__		)		)
              _snake_case		:	List[Any]					     =     len(lowercase__		)
              self.assertEqual(config.output_hidden_states     ,     lowercase__		)
              check_encoder_attentions_output(lowercase__		)

              if self.is_encoder_decoder:
                    _snake_case		:	Union[str, Any]					     =     model_class(lowercase__		)
                    _snake_case		:	List[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"]
              _snake_case		:	str					     =     True
              _snake_case		:	Tuple					     =     model_class(lowercase__		)
              _snake_case		:	int					     =     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
              _snake_case		:	int					     =     True
              _snake_case		:	List[str]					     =     True
              _snake_case		:	Tuple					     =     model_class(lowercase__		)
              _snake_case		:	Optional[Any]					     =     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 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        pass







  def 				UpperCAmelCase_				(		self		)		->						str:






        """simple docstring"""





        pass




def    _a    (							lowerCAmelCase_	):





      """simple docstring"""



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


UpperCAmelCase					:		Dict		=							1E-4





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







  def 				UpperCAmelCase_				(		self		)		->						Dict:






        """simple docstring"""





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

        # change to intended input here
        _snake_case		:	List[str]					     =     _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	Tuple					     =     _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	Tuple					     =     prepare_led_inputs_dict(model.config     ,     lowercase__     ,     lowercase__		)
        _snake_case		:	int					     =     model(**lowercase__		)[0]
        _snake_case		:	Dict					     =     (1, 1_024, 768)
        self.assertEqual(output.shape     ,     lowercase__		)
        # change to expected output here
        _snake_case		:	List[Any]					     =     tf.convert_to_tensor(
            [[2.3_050, 2.8_279, 0.6_531], [-1.8_457, -0.1_455, -3.5_661], [-1.0_186, 0.4_586, -2.2_043]]     ,     )
        tf.debugging.assert_near(output[:, :3, :3]     ,     lowercase__     ,     atol=1E-3		)







  def 				UpperCAmelCase_				(		self		)		->						List[Any]:






        """simple docstring"""





        _snake_case		:	Any					     =     TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''		)

        # change to intended input here
        _snake_case		:	Dict					     =     _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	Dict					     =     _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	List[str]					     =     prepare_led_inputs_dict(model.config     ,     lowercase__     ,     lowercase__		)
        _snake_case		:	Tuple					     =     model(**lowercase__		)[0]
        _snake_case		:	Any					     =     (1, 1_024, model.config.vocab_size)
        self.assertEqual(output.shape     ,     lowercase__		)
        # change to expected output here
        _snake_case		:	Dict					     =     tf.convert_to_tensor(
            [[33.6_507, 6.4_572, 16.8_089], [5.8_739, -2.4_238, 11.2_902], [-3.2_139, -4.3_149, 4.2_783]]     ,     )
        tf.debugging.assert_near(output[:, :3, :3]     ,     lowercase__     ,     atol=1E-3     ,     rtol=1E-3		)


