






"""simple docstring"""



from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


_lowerCAmelCase  :  List[Any]   =			logging.get_logger(__name__)

_lowerCAmelCase  :  List[str]   =			{
    """google/bigbird-pegasus-large-arxiv""": (
        """https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json"""
    ),
    """google/bigbird-pegasus-large-pubmed""": (
        """https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json"""
    ),
    """google/bigbird-pegasus-large-bigpatent""": (
        """https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json"""
    ),
    # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
}




class  lowerCAmelCase__ (						__magic_name__      ):
   SCREAMING_SNAKE_CASE_         ='''bigbird_pegasus'''
   SCREAMING_SNAKE_CASE_         =['''past_key_values''']
   SCREAMING_SNAKE_CASE_         ={
       '''num_attention_heads''': '''encoder_attention_heads''',
       '''hidden_size''': '''d_model''',
       '''attention_probs_dropout_prob''': '''attention_dropout''',
   }
   def __init__(		self			:  str       , snake_case__			:  Any=9_6_1_0_3       , snake_case__			:  Tuple=4_0_9_6       , snake_case__			:  List[Any]=1_6       , snake_case__			:  List[str]=4_0_9_6       , snake_case__			:  Union[str, Any]=1_6       , snake_case__			:  List[str]=1_6       , snake_case__			:  Optional[int]=4_0_9_6       , snake_case__			:  Optional[Any]=1_6       , snake_case__			:  Any=0.0       , snake_case__			:  int=0.0       , snake_case__			:  Optional[Any]=True       , snake_case__			:  Optional[Any]=True       , snake_case__			:  Any="gelu_new"       , snake_case__			:  int=1_0_2_4       , snake_case__			:  Dict=0.1       , snake_case__			:  List[Any]=0.0       , snake_case__			:  str=0.0       , snake_case__			:  str=0.02       , snake_case__			:  Optional[int]=2       , snake_case__			:  Dict=0.0       , snake_case__			:  Tuple=True       , snake_case__			:  Tuple=0       , snake_case__			:  Dict=2       , snake_case__			:  str=1       , snake_case__			:  int="block_sparse"       , snake_case__			:  Optional[int]=6_4       , snake_case__			:  str=3       , snake_case__			:  List[str]=False       , **snake_case__			:  Optional[int]       , ):


        '''simple docstring'''
        UpperCAmelCase__	:						List[Any]       =       vocab_size
        UpperCAmelCase__	:						Tuple       =       max_position_embeddings
        UpperCAmelCase__	:						Dict       =       d_model
        UpperCAmelCase__	:						List[str]       =       encoder_ffn_dim
        UpperCAmelCase__	:						Union[str, Any]       =       encoder_layers
        UpperCAmelCase__	:						Tuple       =       encoder_attention_heads
        UpperCAmelCase__	:						List[str]       =       decoder_ffn_dim
        UpperCAmelCase__	:						List[str]       =       decoder_layers
        UpperCAmelCase__	:						Dict       =       decoder_attention_heads
        UpperCAmelCase__	:						Optional[int]       =       dropout
        UpperCAmelCase__	:						Tuple       =       attention_dropout
        UpperCAmelCase__	:						Optional[int]       =       activation_dropout
        UpperCAmelCase__	:						List[Any]       =       activation_function
        UpperCAmelCase__	:						int       =       init_std
        UpperCAmelCase__	:						List[Any]       =       encoder_layerdrop
        UpperCAmelCase__	:						Tuple       =       decoder_layerdrop
        UpperCAmelCase__	:						Tuple       =       classifier_dropout
        UpperCAmelCase__	:						List[Any]       =       use_cache
        UpperCAmelCase__	:						Any       =       encoder_layers
        UpperCAmelCase__	:						Tuple       =       scale_embedding  # scale factor will be sqrt(d_model) if True

        # extra config
        UpperCAmelCase__	:						List[Any]       =       attention_type
        UpperCAmelCase__	:						List[str]       =       block_size
        UpperCAmelCase__	:						List[str]       =       num_random_blocks
        UpperCAmelCase__	:						str       =       use_bias

        super().__init__(
            pad_token_id=snake_case__       , bos_token_id=snake_case__       , eos_token_id=snake_case__       , is_encoder_decoder=snake_case__       , decoder_start_token_id=snake_case__       , **snake_case__       , )




class  lowerCAmelCase__ (						__magic_name__      ):
   @property
   def 			__a       (		self			:  Optional[int]       ):


        '''simple docstring'''
        if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase__	:						Any       =       OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                 ]       )

             if self.use_past:
                  UpperCAmelCase__	:						Dict       =       {0: "batch"}
                  UpperCAmelCase__	:						Optional[Any]       =       {0: "batch", 1: "past_decoder_sequence + sequence"}
             else:
                  UpperCAmelCase__	:						Optional[int]       =       {0: "batch", 1: "decoder_sequence"}
                  UpperCAmelCase__	:						Union[str, Any]       =       {0: "batch", 1: "decoder_sequence"}

             if self.use_past:
                  self.fill_with_past_key_values_(snake_case__       , direction="inputs"       )
        elif self.task == "causal-lm":
             # TODO: figure this case out.
             UpperCAmelCase__	:						Tuple       =       OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                 ]       )
             if self.use_past:
                  UpperCAmelCase__       ,							UpperCAmelCase__	:						Optional[Any]       =       self.num_layers
                  for i in range(snake_case__       ):
                       UpperCAmelCase__	:						Dict       =       {0: "batch", 2: "past_sequence + sequence"}
                       UpperCAmelCase__	:						Any       =       {0: "batch", 2: "past_sequence + sequence"}
        else:
             UpperCAmelCase__	:						str       =       OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                     ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
                     ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
                 ]       )

        return common_inputs
   @property
   def 			__a       (		self			:  List[Any]       ):


        '''simple docstring'''
        if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase__	:						str       =       super().outputs
        else:
             UpperCAmelCase__	:						Dict       =       super(snake_case__       , self       ).outputs
             if self.use_past:
                  UpperCAmelCase__       ,							UpperCAmelCase__	:						Optional[Any]       =       self.num_layers
                  for i in range(snake_case__       ):
                       UpperCAmelCase__	:						Optional[int]       =       {0: "batch", 2: "past_sequence + sequence"}
                       UpperCAmelCase__	:						Optional[int]       =       {0: "batch", 2: "past_sequence + sequence"}
        return common_outputs
   def 			__a       (		self			:  Union[str, Any]       , snake_case__			:  PreTrainedTokenizer       , snake_case__			:  int = -1       , snake_case__			:  int = -1       , snake_case__			:  bool = False       , snake_case__			:  Optional[TensorType] = None       , ):


        '''simple docstring'''
        UpperCAmelCase__	:						List[Any]       =       self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__       )

        # Generate decoder inputs
        UpperCAmelCase__	:						Optional[int]       =       seq_length if not self.use_past else 1
        UpperCAmelCase__	:						Optional[Any]       =       self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__       )
        UpperCAmelCase__	:						List[str]       =       {f'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
        UpperCAmelCase__	:						List[Any]       =       dict(**snake_case__       , **snake_case__       )

        if self.use_past:
             if not is_torch_available():
                  raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."       )
             else:
                  import torch
             UpperCAmelCase__       ,							UpperCAmelCase__	:						Dict       =       common_inputs["input_ids"].shape
             UpperCAmelCase__	:						Dict       =       common_inputs["decoder_input_ids"].shape[1]
             UpperCAmelCase__       ,							UpperCAmelCase__	:						List[Any]       =       self.num_attention_heads
             UpperCAmelCase__	:						Any       =       (
                 batch,
                 num_encoder_attention_heads,
                 encoder_seq_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             UpperCAmelCase__	:						Optional[int]       =       decoder_seq_length + 3
             UpperCAmelCase__	:						str       =       (
                 batch,
                 num_decoder_attention_heads,
                 decoder_past_length,
                 self._config.hidden_size // num_decoder_attention_heads,
             )

             UpperCAmelCase__	:						Dict       =       torch.cat(
                 [common_inputs["decoder_attention_mask"], torch.ones(snake_case__       , snake_case__       )]       , dim=1       )

             UpperCAmelCase__	:						Tuple       =       []
             # If the number of encoder and decoder layers are present in the model configuration, both are considered
             UpperCAmelCase__       ,							UpperCAmelCase__	:						Any       =       self.num_layers
             UpperCAmelCase__	:						str       =       min(snake_case__       , snake_case__       )
             UpperCAmelCase__	:						Union[str, Any]       =       max(snake_case__       , snake_case__       ) - min_num_layers
             UpperCAmelCase__	:						List[Any]       =       "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

             for _ in range(snake_case__       ):
                  common_inputs["past_key_values"].append(
                      (
                          torch.zeros(snake_case__       ),
                          torch.zeros(snake_case__       ),
                          torch.zeros(snake_case__       ),
                          torch.zeros(snake_case__       ),
                      )       )
             # TODO: test this.
             UpperCAmelCase__	:						List[str]       =       encoder_shape if remaining_side_name == "encoder" else decoder_shape
             for _ in range(snake_case__       , snake_case__       ):
                  common_inputs["past_key_values"].append((torch.zeros(snake_case__       ), torch.zeros(snake_case__       ))       )
        return common_inputs
   def 			__a       (		self			:  Dict       , snake_case__			:  PreTrainedTokenizer       , snake_case__			:  int = -1       , snake_case__			:  int = -1       , snake_case__			:  bool = False       , snake_case__			:  Optional[TensorType] = None       , ):


        '''simple docstring'''
        UpperCAmelCase__	:						List[str]       =       self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__       )

        if self.use_past:
             if not is_torch_available():
                  raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."       )
             else:
                  import torch
             UpperCAmelCase__       ,							UpperCAmelCase__	:						str       =       common_inputs["input_ids"].shape
             # Not using the same length for past_key_values
             UpperCAmelCase__	:						Optional[Any]       =       seqlen + 2
             UpperCAmelCase__       ,							UpperCAmelCase__	:						List[str]       =       self.num_layers
             UpperCAmelCase__       ,							UpperCAmelCase__	:						List[str]       =       self.num_attention_heads
             UpperCAmelCase__	:						str       =       (
                 batch,
                 num_encoder_attention_heads,
                 past_key_values_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )

             UpperCAmelCase__	:						Any       =       common_inputs["attention_mask"].dtype
             UpperCAmelCase__	:						Union[str, Any]       =       torch.cat(
                 [common_inputs["attention_mask"], torch.ones(snake_case__       , snake_case__       , dtype=snake_case__       )]       , dim=1       )
             UpperCAmelCase__	:						Optional[int]       =       [
                 (torch.zeros(snake_case__       ), torch.zeros(snake_case__       )) for _ in range(snake_case__       )
             ]
        return common_inputs
   def 			__a       (		self			:  Dict       , snake_case__			:  PreTrainedTokenizer       , snake_case__			:  int = -1       , snake_case__			:  int = -1       , snake_case__			:  bool = False       , snake_case__			:  Optional[TensorType] = None       , ):


        '''simple docstring'''
        # Copied from OnnxConfig.generate_dummy_inputs
        # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
        # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
        UpperCAmelCase__	:						Any       =       compute_effective_axis_dimension(
            snake_case__       , fixed_dimension=OnnxConfig.default_fixed_batch       , num_token_to_add=0       )

        # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
        UpperCAmelCase__	:						Optional[int]       =       tokenizer.num_special_tokens_to_add(snake_case__       )
        UpperCAmelCase__	:						str       =       compute_effective_axis_dimension(
            snake_case__       , fixed_dimension=OnnxConfig.default_fixed_sequence       , num_token_to_add=snake_case__       )

        # Generate dummy inputs according to compute batch and sequence
        UpperCAmelCase__	:						List[Any]       =       [" ".join([tokenizer.unk_token]       ) * seq_length] * batch_size
        UpperCAmelCase__	:						str       =       dict(tokenizer(snake_case__       , return_tensors=snake_case__       )       )
        return common_inputs
   def 			__a       (		self			:  Union[str, Any]       , snake_case__			:  PreTrainedTokenizer       , snake_case__			:  int = -1       , snake_case__			:  int = -1       , snake_case__			:  bool = False       , snake_case__			:  Optional[TensorType] = None       , ):


        '''simple docstring'''
        if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase__	:						Union[str, Any]       =       self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                 snake_case__       , batch_size=snake_case__       , seq_length=snake_case__       , is_pair=snake_case__       , framework=snake_case__       )

        elif self.task == "causal-lm":
             UpperCAmelCase__	:						List[Any]       =       self._generate_dummy_inputs_for_causal_lm(
                 snake_case__       , batch_size=snake_case__       , seq_length=snake_case__       , is_pair=snake_case__       , framework=snake_case__       )
        else:
             UpperCAmelCase__	:						Optional[Any]       =       self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                 snake_case__       , batch_size=snake_case__       , seq_length=snake_case__       , is_pair=snake_case__       , framework=snake_case__       )

        return common_inputs




   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  int       , snake_case__			:  str       , snake_case__			:  Tuple       , snake_case__			:  Optional[int]       ):


        '''simple docstring'''
        if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase__	:						Tuple       =       super()._flatten_past_key_values_(snake_case__       , snake_case__       , snake_case__       , snake_case__       )
        else:
             UpperCAmelCase__	:						Union[str, Any]       =       super(snake_case__       , self       )._flatten_past_key_values_(
                 snake_case__       , snake_case__       , snake_case__       , snake_case__       )










"""simple docstring"""



#
# Sort large text files in a minimum amount of memory
#
import argparse
import os




class  lowerCAmelCase__ :
   SCREAMING_SNAKE_CASE_         ='''block_{0}.dat'''
   def __init__(		self			:  Tuple       , snake_case__			:  List[str]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						List[Any]       =       filename
        UpperCAmelCase__	:						List[Any]       =       []
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  Union[str, Any]       , snake_case__			:  Union[str, Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Dict       =       self.BLOCK_FILENAME_FORMAT.format(snake_case__       )
        with open(snake_case__       , "w"       ) as file:
             file.write(snake_case__       )
        self.block_filenames.append(snake_case__       )
   def 			__a       (		self			:  List[str]       ):


        '''simple docstring'''
        return self.block_filenames
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  List[str]       , snake_case__			:  str=None       ):


        '''simple docstring'''
        UpperCAmelCase__	:						int       =       0
        with open(self.filename       ) as file:
             while True:
                  UpperCAmelCase__	:						Union[str, Any]       =       file.readlines(snake_case__       )

                  if lines == []:
                       break

                  if sort_key is None:
                       lines.sort()
                  else:
                       lines.sort(key=snake_case__       )

                  self.write_block("".join(snake_case__       )       , snake_case__       )
                  i += 1




   def 			__a       (		self			:  Optional[Any]       ):


        '''simple docstring'''
        map(os.remove       , self.block_filenames       )




class  lowerCAmelCase__ :
   def 			__a       (		self			:  Union[str, Any]       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Any       =       -1
        UpperCAmelCase__	:						Tuple       =       None

        for i in range(len(snake_case__       )       ):
             if min_str is None or choices[i] < min_str:
                  UpperCAmelCase__	:						List[str]       =       i

        return min_index




class  lowerCAmelCase__ :
   def __init__(		self			:  Dict       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Optional[Any]       =       files
        UpperCAmelCase__	:						str       =       set()
        UpperCAmelCase__	:						Any       =       len(snake_case__       )
        UpperCAmelCase__	:						Tuple       =       {i: None for i in range(self.num_buffers       )}
   def 			__a       (		self			:  List[str]       ):


        '''simple docstring'''
        return {
            i: self.buffers[i] for i in range(self.num_buffers       ) if i not in self.empty
        }
   def 			__a       (		self			:  str       ):


        '''simple docstring'''
        for i in range(self.num_buffers       ):
             if self.buffers[i] is None and i not in self.empty:
                  UpperCAmelCase__	:						Tuple       =       self.files[i].readline()

                  if self.buffers[i] == "":
                       self.empty.add(snake_case__       )
                       self.files[i].close()

        if len(self.empty       ) == self.num_buffers:
             return False

        return True




   def 			__a       (		self			:  List[str]       , snake_case__			:  Tuple       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Any       =       self.buffers[index]
        UpperCAmelCase__	:						Union[str, Any]       =       None

        return value




class  lowerCAmelCase__ :
   def __init__(		self			:  Optional[Any]       , snake_case__			:  int       ):


        '''simple docstring'''
        UpperCAmelCase__	:						str       =       merge_strategy
   def 			__a       (		self			:  Dict       , snake_case__			:  Dict       , snake_case__			:  Optional[int]       , snake_case__			:  Optional[int]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Dict       =       FilesArray(self.get_file_handles(snake_case__       , snake_case__       )       )
        with open(snake_case__       , "w"       , snake_case__       ) as outfile:
             while buffers.refresh():
                  UpperCAmelCase__	:						Any       =       self.merge_strategy.select(buffers.get_dict()       )
                  outfile.write(buffers.unshift(snake_case__       )       )




   def 			__a       (		self			:  Tuple       , snake_case__			:  Dict       , snake_case__			:  List[str]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						str       =       {}

        for i in range(len(snake_case__       )       ):
             UpperCAmelCase__	:						List[Any]       =       open(filenames[i]       , "r"       , snake_case__       )  # noqa: UP015

        return files




class  lowerCAmelCase__ :
   def __init__(		self			:  str       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						List[str]       =       block_size
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  Optional[int]       , snake_case__			:  Any=None       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Union[str, Any]       =       self.get_number_blocks(snake_case__       , self.block_size       )
        UpperCAmelCase__	:						Union[str, Any]       =       FileSplitter(snake_case__       )
        splitter.split(self.block_size       , snake_case__       )

        UpperCAmelCase__	:						int       =       FileMerger(NWayMerge()       )
        UpperCAmelCase__	:						int       =       self.block_size / (num_blocks + 1)
        merger.merge(splitter.get_block_filenames()       , filename + ".out"       , snake_case__       )

        splitter.cleanup()




   def 			__a       (		self			:  Union[str, Any]       , snake_case__			:  Optional[int]       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        return (os.stat(snake_case__       ).st_size / block_size) + 1




def   SCREAMING_SNAKE_CASE__					(     snake_case					:   Tuple							)->		Union[str, Any]:






     '''simple docstring'''



     if string[-1].lower() == "k":
          return int(string[:-1]							) * 1024
     elif string[-1].lower() == "m":
          return int(string[:-1]							) * 1024 * 1024
     elif string[-1].lower() == "g":
          return int(string[:-1]							) * 1024 * 1024 * 1024
     else:
          return int(snake_case							)


def   SCREAMING_SNAKE_CASE__					(     )->		Any:






     '''simple docstring'''



     UpperCAmelCase__	:						Any       =       argparse.ArgumentParser()
     parser.add_argument(
         "-m" ,				"--mem" ,				help="amount of memory to use for sorting" ,				default="100M"							)
     parser.add_argument(
         "filename" ,				metavar="<filename>" ,				nargs=1 ,				help="name of file to sort"							)
     UpperCAmelCase__	:						int       =       parser.parse_args()

     UpperCAmelCase__	:						Optional[int]       =       ExternalSort(parse_memory(args.mem							)							)
     sorter.sort(args.filename[0]							)


if __name__ == "__main__":
 main()


