






from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin


class 			SCREAMING_SNAKE_CASE      (      lowerCAmelCase   ,		lowerCAmelCase   ,		lowerCAmelCase       ):
     '''simple docstring'''


     UpperCamelCase_				:		Optional[int]								=						[r'''h\.\d+\.attn\.bias''', r'''h\.\d+\.attn\.masked_bias''']


     @register_to_config
     def __init__(					self							:						Union[str, Any]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						Optional[int] = None       ,			UpperCAmelCase_							:						int = 5_0257       ,			UpperCAmelCase_							:						int = 1024       ,			UpperCAmelCase_							:						int = 768       ,			UpperCAmelCase_							:						int = 12       ,			UpperCAmelCase_							:						int = 12       ,			UpperCAmelCase_							:						Optional[int] = None       ,			UpperCAmelCase_							:						str = "gelu_new"       ,			UpperCAmelCase_							:						float = 0.1       ,			UpperCAmelCase_							:						float = 0.1       ,			UpperCAmelCase_							:						float = 0.1       ,			UpperCAmelCase_							:						float = 1E-5       ,			UpperCAmelCase_							:						float = 0.02       ,			UpperCAmelCase_							:						bool = True       ,			UpperCAmelCase_							:						bool = True       ,			UpperCAmelCase_							:						bool = False       ,			UpperCAmelCase_							:						bool = False       ,			):
       super().__init__()

       SCREAMING_SNAKE_CASE							:   Tuple						    =     prefix_length

       if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
         raise ValueError(
             f'''`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'''
             f''' `n_embd`: {n_embd} are not equal.'''   )

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     prefix_inner_dim
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     prefix_hidden_dim

       SCREAMING_SNAKE_CASE							:   Dict						    =     (
           nn.Linear(self.prefix_inner_dim       ,			self.prefix_hidden_dim   )
           if self.prefix_hidden_dim is not None
           else nn.Identity()
       )
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     (
           nn.Linear(self.prefix_hidden_dim       ,			UpperCAmelCase_   ) if self.prefix_hidden_dim is not None else nn.Identity()
       )

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     GPTaConfig(
           vocab_size=UpperCAmelCase_       ,			n_positions=UpperCAmelCase_       ,			n_embd=UpperCAmelCase_       ,			n_layer=UpperCAmelCase_       ,			n_head=UpperCAmelCase_       ,			n_inner=UpperCAmelCase_       ,			activation_function=UpperCAmelCase_       ,			resid_pdrop=UpperCAmelCase_       ,			embd_pdrop=UpperCAmelCase_       ,			attn_pdrop=UpperCAmelCase_       ,			layer_norm_epsilon=UpperCAmelCase_       ,			initializer_range=UpperCAmelCase_       ,			scale_attn_weights=UpperCAmelCase_       ,			use_cache=UpperCAmelCase_       ,			scale_attn_by_inverse_layer_idx=UpperCAmelCase_       ,			reorder_and_upcast_attn=UpperCAmelCase_       ,			)
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     GPTaLMHeadModel(UpperCAmelCase_   )


     def 					_A							(					self							:						Optional[Any]       ,			UpperCAmelCase_							:						torch.Tensor       ,			UpperCAmelCase_							:						torch.Tensor       ,			UpperCAmelCase_							:						Optional[torch.Tensor] = None       ,			UpperCAmelCase_							:						Optional[torch.Tensor] = None       ,			):
       SCREAMING_SNAKE_CASE							:   List[str]						    =     self.transformer.transformer.wte(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   int						    =     self.encode_prefix(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     self.decode_prefix(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     torch.cat((prefix_embeds, embedding_text)       ,			dim=1   )

       if labels is not None:
         SCREAMING_SNAKE_CASE							:   Any						    =     self.get_dummy_token(input_ids.shape[0]       ,			input_ids.device   )
         SCREAMING_SNAKE_CASE							:   Any						    =     torch.cat((dummy_token, input_ids)       ,			dim=1   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     self.transformer(inputs_embeds=UpperCAmelCase_       ,			labels=UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_   )
       if self.prefix_hidden_dim is not None:
         return out, hidden
       else:
         return out


     def 					_A							(					self							:						List[Any]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						torch.device   ):
       return torch.zeros(UpperCAmelCase_       ,			self.prefix_length       ,			dtype=torch.intaa       ,			device=UpperCAmelCase_   )


     def 					_A							(					self							:						Optional[int]       ,			UpperCAmelCase_							:						Union[str, Any]   ):
       return self.encode_prefix(UpperCAmelCase_   )


     @torch.no_grad()
     def 					_A							(					self							:						List[Any]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						List[Any]   ):
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     torch.split(UpperCAmelCase_       ,			1       ,			dim=0   )
       SCREAMING_SNAKE_CASE							:   List[str]						    =     []
       SCREAMING_SNAKE_CASE							:   Dict						    =     []
       for feature in features:
         SCREAMING_SNAKE_CASE							:   List[Any]						    =     self.decode_prefix(feature.to(UpperCAmelCase_   )   )  # back to the clip feature
         # Only support beam search for now
         SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Dict						    =     self.generate_beam(
             input_embeds=UpperCAmelCase_       ,			device=UpperCAmelCase_       ,			eos_token_id=UpperCAmelCase_   )
         generated_tokens.append(output_tokens[0]   )
         generated_seq_lengths.append(seq_lengths[0]   )
       SCREAMING_SNAKE_CASE							:   int						    =     torch.stack(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Dict						    =     torch.stack(UpperCAmelCase_   )
       return generated_tokens, generated_seq_lengths






     @torch.no_grad()
     def 					_A							(					self							:						Any       ,			UpperCAmelCase_							:						Dict=None       ,			UpperCAmelCase_							:						Union[str, Any]=None       ,			UpperCAmelCase_							:						Optional[int]=None       ,			UpperCAmelCase_							:						int = 5       ,			UpperCAmelCase_							:						int = 67       ,			UpperCAmelCase_							:						float = 1.0       ,			UpperCAmelCase_							:						Optional[int] = None       ,			):
       SCREAMING_SNAKE_CASE							:   Tuple						    =     eos_token_id
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     None
       SCREAMING_SNAKE_CASE							:   int						    =     None
       SCREAMING_SNAKE_CASE							:   str						    =     torch.ones(UpperCAmelCase_       ,			device=UpperCAmelCase_       ,			dtype=torch.int   )
       SCREAMING_SNAKE_CASE							:   Any						    =     torch.zeros(UpperCAmelCase_       ,			device=UpperCAmelCase_       ,			dtype=torch.bool   )

       if input_embeds is not None:
         SCREAMING_SNAKE_CASE							:   Tuple						    =     input_embeds
       else:
         SCREAMING_SNAKE_CASE							:   int						    =     self.transformer.transformer.wte(UpperCAmelCase_   )

       for i in range(UpperCAmelCase_   ):
         SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     self.transformer(inputs_embeds=UpperCAmelCase_   )
         SCREAMING_SNAKE_CASE							:   Tuple						    =     outputs.logits
         SCREAMING_SNAKE_CASE							:   int						    =     logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
         SCREAMING_SNAKE_CASE							:   Dict						    =     logits.softmax(-1   ).log()

         if scores is None:
           SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Any						    =     logits.topk(UpperCAmelCase_       ,			-1   )
           SCREAMING_SNAKE_CASE							:   Tuple						    =     generated.expand(UpperCAmelCase_       ,			*generated.shape[1:]   )
           SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Dict						    =     next_tokens.permute(1       ,			0   ), scores.squeeze(0   )
           if tokens is None:
             SCREAMING_SNAKE_CASE							:   str						    =     next_tokens
           else:
             SCREAMING_SNAKE_CASE							:   Optional[int]						    =     tokens.expand(UpperCAmelCase_       ,			*tokens.shape[1:]   )
             SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     torch.cat((tokens, next_tokens)       ,			dim=1   )
         else:
           SCREAMING_SNAKE_CASE							:   Tuple						    =     -float(np.inf   )
           SCREAMING_SNAKE_CASE							:   List[str]						    =     0
           SCREAMING_SNAKE_CASE							:   Dict						    =     scores[:, None] + logits
           seq_lengths[~is_stopped] += 1
           SCREAMING_SNAKE_CASE							:   Any						    =     scores_sum / seq_lengths[:, None]
           SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Tuple						    =     scores_sum_average.view(-1   ).topk(UpperCAmelCase_       ,			-1   )
           SCREAMING_SNAKE_CASE							:   str						    =     next_tokens // scores_sum.shape[1]
           SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     seq_lengths[next_tokens_source]
           SCREAMING_SNAKE_CASE							:   List[str]						    =     next_tokens % scores_sum.shape[1]
           SCREAMING_SNAKE_CASE							:   Dict						    =     next_tokens.unsqueeze(1   )
           SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     tokens[next_tokens_source]
           SCREAMING_SNAKE_CASE							:   Optional[int]						    =     torch.cat((tokens, next_tokens)       ,			dim=1   )
           SCREAMING_SNAKE_CASE							:   Any						    =     generated[next_tokens_source]
           SCREAMING_SNAKE_CASE							:   Optional[int]						    =     scores_sum_average * seq_lengths
           SCREAMING_SNAKE_CASE							:   Tuple						    =     is_stopped[next_tokens_source]

         SCREAMING_SNAKE_CASE							:   int						    =     self.transformer.transformer.wte(next_tokens.squeeze()   ).view(generated.shape[0]       ,			1       ,			-1   )
         SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     torch.cat((generated, next_token_embed)       ,			dim=1   )
         SCREAMING_SNAKE_CASE							:   Tuple						    =     is_stopped + next_tokens.eq(UpperCAmelCase_   ).squeeze()
         if is_stopped.all():
           break

       SCREAMING_SNAKE_CASE							:   List[str]						    =     scores / seq_lengths
       SCREAMING_SNAKE_CASE							:   Tuple						    =     scores.argsort(descending=UpperCAmelCase_   )
       # tokens tensors are already padded to max_seq_length
       SCREAMING_SNAKE_CASE							:   int						    =     [tokens[i] for i in order]
       SCREAMING_SNAKE_CASE							:   List[str]						    =     torch.stack(UpperCAmelCase_       ,			dim=0   )
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     torch.tensor([seq_lengths[i] for i in order]       ,			dtype=seq_lengths.dtype   )
       return output_texts, seq_lengths












from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import TensorType, logging


if TYPE_CHECKING:
						from ...onnx.config import PatchingSpec
						from ...tokenization_utils_base import PreTrainedTokenizerBase


SCREAMING_SNAKE_CASE:     str									= logging.get_logger(__name__)

SCREAMING_SNAKE_CASE:     Optional[int]									= {
    '''allenai/longformer-base-4096''': '''https://huggingface.co/allenai/longformer-base-4096/resolve/main/config.json''',
    '''allenai/longformer-large-4096''': '''https://huggingface.co/allenai/longformer-large-4096/resolve/main/config.json''',
    '''allenai/longformer-large-4096-finetuned-triviaqa''': (
        '''https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/config.json'''
    ),
    '''allenai/longformer-base-4096-extra.pos.embd.only''': (
        '''https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/config.json'''
    ),
    '''allenai/longformer-large-4096-extra.pos.embd.only''': (
        '''https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/config.json'''
    ),
}



class    lowercase_		(SCREAMING_SNAKE_CASE__ ):
						lowerCAmelCase__			      ="longformer"



						def __init__(			self	:   Optional[Any]					,							snake_case__	:   Union[List[int], int] = 5_12					,							snake_case__	:   int = 2					,							snake_case__	:   int = 1					,							snake_case__	:   int = 0					,							snake_case__	:   int = 2					,							snake_case__	:   int = 3_05_22					,							snake_case__	:   int = 7_68					,							snake_case__	:   int = 12					,							snake_case__	:   int = 12					,							snake_case__	:   int = 30_72					,							snake_case__	:   str = "gelu"					,							snake_case__	:   float = 0.1					,							snake_case__	:   float = 0.1					,							snake_case__	:   int = 5_12					,							snake_case__	:   int = 2					,							snake_case__	:   float = 0.02					,							snake_case__	:   float = 1e-12					,							snake_case__	:   bool = False					,							**snake_case__	:   Tuple					,							):

								"""simple docstring"""

								super().__init__(pad_token_id=snake_case__					,							**snake_case__  )

								SCREAMING_SNAKE_CASE_					      =							attention_window
								SCREAMING_SNAKE_CASE_					      =							sep_token_id
								SCREAMING_SNAKE_CASE_					      =							bos_token_id
								SCREAMING_SNAKE_CASE_					      =							eos_token_id
								SCREAMING_SNAKE_CASE_					      =							vocab_size
								SCREAMING_SNAKE_CASE_					      =							hidden_size
								SCREAMING_SNAKE_CASE_					      =							num_hidden_layers
								SCREAMING_SNAKE_CASE_					      =							num_attention_heads
								SCREAMING_SNAKE_CASE_					      =							hidden_act
								SCREAMING_SNAKE_CASE_					      =							intermediate_size
								SCREAMING_SNAKE_CASE_					      =							hidden_dropout_prob
								SCREAMING_SNAKE_CASE_					      =							attention_probs_dropout_prob
								SCREAMING_SNAKE_CASE_					      =							max_position_embeddings
								SCREAMING_SNAKE_CASE_					      =							type_vocab_size
								SCREAMING_SNAKE_CASE_					      =							initializer_range
								SCREAMING_SNAKE_CASE_					      =							layer_norm_eps
								SCREAMING_SNAKE_CASE_					      =							onnx_export



class    lowercase_		(SCREAMING_SNAKE_CASE__ ):



						def __init__(			self	:   Dict					,							snake_case__	:   "PretrainedConfig"					,							snake_case__	:   str = "default"					,							snake_case__	:   "List[PatchingSpec]" = None  ):

								"""simple docstring"""

								super().__init__(snake_case__					,							snake_case__					,							snake_case__  )
								SCREAMING_SNAKE_CASE_					      =							True



						@property
						def 	__a					(			self	:   Union[str, Any]  ):

								"""simple docstring"""

								if self.task == "multiple-choice":
										SCREAMING_SNAKE_CASE_					      =							{0: 'batch', 1: 'choice', 2: 'sequence'}
								else:
										SCREAMING_SNAKE_CASE_					      =							{0: 'batch', 1: 'sequence'}
								return OrderedDict(
								    [
								        ('input_ids', dynamic_axis),
								        ('attention_mask', dynamic_axis),
								        ('global_attention_mask', dynamic_axis),
								    ]  )



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

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							super().outputs
								if self.task == "default":
										SCREAMING_SNAKE_CASE_					      =							{0: 'batch'}
								return outputs



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

								"""simple docstring"""

								return 1e-4



						@property
						def 	__a					(			self	:   List[str]  ):

								"""simple docstring"""

								return max(super().default_onnx_opset					,							14  )






						def 	__a					(			self	:   List[str]					,							snake_case__	:   "PreTrainedTokenizerBase"					,							snake_case__	:   int = -1					,							snake_case__	:   int = -1					,							snake_case__	:   bool = False					,							snake_case__	:   Optional[TensorType] = None					,							):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							super().generate_dummy_inputs(
								    preprocessor=snake_case__					,							batch_size=snake_case__					,							seq_length=snake_case__					,							is_pair=snake_case__					,							framework=snake_case__  )
								import torch

								# for some reason, replacing this code by inputs["global_attention_mask"] = torch.randint(2, inputs["input_ids"].shape, dtype=torch.int64)
								# makes the export fail randomly
								SCREAMING_SNAKE_CASE_					      =							torch.zeros_like(inputs['input_ids']  )
								# make every second token global
								SCREAMING_SNAKE_CASE_					      =							1

								return inputs