




"""simple docstring"""






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   a__     (       lowercase_					,	lowercase_					,	lowercase_  ):
						snake_case__             = [r'''h\.\d+\.attn\.bias''', r'''h\.\d+\.attn\.masked_bias''']





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




											"""simple docstring"""







											super().__init__()

											_lowerCAmelCase:str									=  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.')

											_lowerCAmelCase:Optional[int]									=  prefix_inner_dim
											_lowerCAmelCase:Union[str, Any]									=  prefix_hidden_dim

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

											_lowerCAmelCase:Any									=  GPTaConfig(
											    vocab_size=lowerCamelCase_     ,n_positions=lowerCamelCase_     ,n_embd=lowerCamelCase_     ,n_layer=lowerCamelCase_     ,n_head=lowerCamelCase_     ,n_inner=lowerCamelCase_     ,activation_function=lowerCamelCase_     ,resid_pdrop=lowerCamelCase_     ,embd_pdrop=lowerCamelCase_     ,attn_pdrop=lowerCamelCase_     ,layer_norm_epsilon=lowerCamelCase_     ,initializer_range=lowerCamelCase_     ,scale_attn_weights=lowerCamelCase_     ,use_cache=lowerCamelCase_     ,scale_attn_by_inverse_layer_idx=lowerCamelCase_     ,reorder_and_upcast_attn=lowerCamelCase_     ,)
											_lowerCAmelCase:int									=  GPTaLMHeadModel(lowerCamelCase_)





						def       __UpperCamelCase     (							self      :	Optional[int]     ,a__      :	torch.Tensor     ,a__      :	torch.Tensor     ,a__      :	Optional[torch.Tensor] = None     ,a__      :	Optional[torch.Tensor] = None     ,)						->       Dict:




											"""simple docstring"""







											_lowerCAmelCase:int									=  self.transformer.transformer.wte(lowerCamelCase_)
											_lowerCAmelCase:int									=  self.encode_prefix(lowerCamelCase_)
											_lowerCAmelCase:List[Any]									=  self.decode_prefix(lowerCamelCase_)
											_lowerCAmelCase:Any									=  torch.cat((prefix_embeds, embedding_text)     ,dim=1)

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





						def       __UpperCamelCase     (							self      :	Optional[Any]     ,a__      :	int     ,a__      :	torch.device)						->       str:




											"""simple docstring"""







											return torch.zeros(lowerCamelCase_     ,self.prefix_length     ,dtype=torch.intaa     ,device=lowerCamelCase_)





						def       __UpperCamelCase     (							self      :	List[Any]     ,a__      :	Union[str, Any])						->       Optional[int]:




											"""simple docstring"""







											return self.encode_prefix(lowerCamelCase_)





						@torch.no_grad()
						def       __UpperCamelCase     (							self      :	List[str]     ,a__      :	List[str]     ,a__      :	Optional[Any]     ,a__      :	List[Any])						->       Any:




											"""simple docstring"""







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







						@torch.no_grad()
						def       __UpperCamelCase     (							self      :	int     ,a__      :	List[Any]=None     ,a__      :	List[Any]=None     ,a__      :	List[str]=None     ,a__      :	int = 5     ,a__      :	int = 67     ,a__      :	float = 1.0     ,a__      :	Optional[int] = None     ,)						->       str:




											"""simple docstring"""







											_lowerCAmelCase:str									=  eos_token_id
											_lowerCAmelCase:int									=  None
											_lowerCAmelCase:List[str]									=  None
											_lowerCAmelCase:Tuple									=  torch.ones(lowerCamelCase_     ,device=lowerCamelCase_     ,dtype=torch.int)
											_lowerCAmelCase:Tuple									=  torch.zeros(lowerCamelCase_     ,device=lowerCamelCase_     ,dtype=torch.bool)

											if input_embeds is not None:
																_lowerCAmelCase:Any									=  input_embeds
											else:
																_lowerCAmelCase:Optional[int]									=  self.transformer.transformer.wte(lowerCamelCase_)

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

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

																_lowerCAmelCase:List[str]									=  self.transformer.transformer.wte(next_tokens.squeeze()).view(generated.shape[0]     ,1     ,-1)
																_lowerCAmelCase:Tuple									=  torch.cat((generated, next_token_embed)     ,dim=1)
																_lowerCAmelCase:Tuple									=  is_stopped + next_tokens.eq(lowerCamelCase_).squeeze()
																if is_stopped.all():
																					break

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












"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__						=			{
    '''configuration_convbert''': ['''CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ConvBertConfig''', '''ConvBertOnnxConfig'''],
    '''tokenization_convbert''': ['''ConvBertTokenizer'''],
}

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase__						=			['''ConvBertTokenizerFast''']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase__						=			[
				    '''CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
				    '''ConvBertForMaskedLM''',
				    '''ConvBertForMultipleChoice''',
				    '''ConvBertForQuestionAnswering''',
				    '''ConvBertForSequenceClassification''',
				    '''ConvBertForTokenClassification''',
				    '''ConvBertLayer''',
				    '''ConvBertModel''',
				    '''ConvBertPreTrainedModel''',
				    '''load_tf_weights_in_convbert''',
				]


try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase__						=			[
				    '''TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
				    '''TFConvBertForMaskedLM''',
				    '''TFConvBertForMultipleChoice''',
				    '''TFConvBertForQuestionAnswering''',
				    '''TFConvBertForSequenceClassification''',
				    '''TFConvBertForTokenClassification''',
				    '''TFConvBertLayer''',
				    '''TFConvBertModel''',
				    '''TFConvBertPreTrainedModel''',
				]


if TYPE_CHECKING:
				from .configuration_convbert import CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ConvBertConfig, ConvBertOnnxConfig
				from .tokenization_convbert import ConvBertTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_convbert_fast import ConvBertTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_convbert import (
								    CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    ConvBertForMaskedLM,
								    ConvBertForMultipleChoice,
								    ConvBertForQuestionAnswering,
								    ConvBertForSequenceClassification,
								    ConvBertForTokenClassification,
								    ConvBertLayer,
								    ConvBertModel,
								    ConvBertPreTrainedModel,
								    load_tf_weights_in_convbert,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_convbert import (
								    TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFConvBertForMaskedLM,
								    TFConvBertForMultipleChoice,
								    TFConvBertForQuestionAnswering,
								    TFConvBertForSequenceClassification,
								    TFConvBertForTokenClassification,
								    TFConvBertLayer,
								    TFConvBertModel,
								    TFConvBertPreTrainedModel,
								)


else:
				import sys

				UpperCamelCase__						=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






