




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      (   __UpperCamelCase	,	__UpperCamelCase	,	__UpperCamelCase				):
 __lowercase			:							Optional[Any]       					=			[r'h\.\d+\.attn\.bias', r'h\.\d+\.attn\.masked_bias']
 @register_to_config
 def __init__(       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None    ,			lowerCAmelCase__ = 50_257    ,			lowerCAmelCase__ = 1_024    ,			lowerCAmelCase__ = 768    ,			lowerCAmelCase__ = 12    ,			lowerCAmelCase__ = 12    ,			lowerCAmelCase__ = None    ,			lowerCAmelCase__ = "gelu_new"    ,			lowerCAmelCase__ = 0.1    ,			lowerCAmelCase__ = 0.1    ,			lowerCAmelCase__ = 0.1    ,			lowerCAmelCase__ = 1E-5    ,			lowerCAmelCase__ = 0.0_2    ,			lowerCAmelCase__ = True    ,			lowerCAmelCase__ = True    ,			lowerCAmelCase__ = False    ,			lowerCAmelCase__ = False    ,			) ->    Optional[int]:




        '''simple docstring'''
        super().__init__()

        lowercase__:       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.'			)

        lowercase__:       List[str]        =       prefix_inner_dim
        lowercase__:       Optional[int]        =       prefix_hidden_dim

        lowercase__:       Optional[int]        =       (
            nn.Linear(self.prefix_inner_dim    ,			self.prefix_hidden_dim			)
            if self.prefix_hidden_dim is not None
            else nn.Identity()
        )
        lowercase__:       Any        =       (
            nn.Linear(self.prefix_hidden_dim    ,			lowerCAmelCase__			) if self.prefix_hidden_dim is not None else nn.Identity()
        )

        lowercase__:       str        =       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__    ,			)
        lowercase__:       str        =       GPTaLMHeadModel(lowerCAmelCase__			)
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None    ,			lowerCAmelCase__ = None    ,			) ->    List[str]:




        '''simple docstring'''
        lowercase__:       List[Any]        =       self.transformer.transformer.wte(lowerCAmelCase__			)
        lowercase__:       Any        =       self.encode_prefix(lowerCAmelCase__			)
        lowercase__:       Union[str, Any]        =       self.decode_prefix(lowerCAmelCase__			)
        lowercase__:       Tuple        =       torch.cat((prefix_embeds, embedding_text)    ,			dim=1			)

        if labels is not None:
               lowercase__:       Any        =       self.get_dummy_token(input_ids.shape[0]    ,			input_ids.device			)
               lowercase__:       Tuple        =       torch.cat((dummy_token, input_ids)    ,			dim=1			)
        lowercase__:       Tuple        =       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 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__			) ->    torch.Tensor:




        '''simple docstring'''
        return torch.zeros(lowerCAmelCase__    ,			self.prefix_length    ,			dtype=torch.intaa    ,			device=lowerCAmelCase__			)
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__			) ->    str:




        '''simple docstring'''
        return self.encode_prefix(lowerCAmelCase__			)
 @torch.no_grad()
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__    ,			lowerCAmelCase__			) ->    Any:




        '''simple docstring'''
        lowercase__:       Any        =       torch.split(lowerCAmelCase__    ,			1    ,			dim=0			)
        lowercase__:       Optional[Any]        =       []
        lowercase__:       Union[str, Any]        =       []
        for feature in features:
               lowercase__:       int        =       self.decode_prefix(feature.to(lowerCAmelCase__			)			)  # back to the clip feature
               # Only support beam search for now
               lowercase__:       str        =       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]			)
        lowercase__:       Any        =       torch.stack(lowerCAmelCase__			)
        lowercase__:       Dict        =       torch.stack(lowerCAmelCase__			)
        return generated_tokens, generated_seq_lengths
 @torch.no_grad()
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__=None    ,			lowerCAmelCase__=None    ,			lowerCAmelCase__=None    ,			lowerCAmelCase__ = 5    ,			lowerCAmelCase__ = 67    ,			lowerCAmelCase__ = 1.0    ,			lowerCAmelCase__ = None    ,			) ->    int:




        '''simple docstring'''
        lowercase__:       Optional[Any]        =       eos_token_id
        lowercase__:       int        =       None
        lowercase__:       int        =       None
        lowercase__:       Optional[int]        =       torch.ones(lowerCAmelCase__    ,			device=lowerCAmelCase__    ,			dtype=torch.int			)
        lowercase__:       Any        =       torch.zeros(lowerCAmelCase__    ,			device=lowerCAmelCase__    ,			dtype=torch.bool			)

        if input_embeds is not None:
               lowercase__:       Union[str, Any]        =       input_embeds
        else:
               lowercase__:       Any        =       self.transformer.transformer.wte(lowerCAmelCase__			)

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

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

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

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






import os
from typing import Dict, List, Union

import tensorflow as tf
from keras_nlp.tokenizers import BytePairTokenizer
from tensorflow_text import pad_model_inputs

from .tokenization_gpta import GPTaTokenizer





class   __a      (   tf.keras.layers.Layer				):
 def __init__(       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None    ,			lowerCAmelCase__ = None			) ->    int:




        '''simple docstring'''
        super().__init__()
        lowercase__:       Union[str, Any]        =       pad_token_id
        lowercase__:       List[str]        =       max_length
        lowercase__:       int        =       vocab
        lowercase__:       List[Any]        =       merges
        lowercase__:       str        =       BytePairTokenizer(lowerCAmelCase__    ,			lowerCAmelCase__    ,			sequence_length=lowerCAmelCase__			)
 @classmethod
 def 		SCREAMING_SNAKE_CASE__       (       cls    ,			lowerCAmelCase__    ,			*lowerCAmelCase__    ,			**lowerCAmelCase__			) ->    Any:




        '''simple docstring'''
        lowercase__:       Tuple        =       [' '.join(lowerCAmelCase__			) for m in tokenizer.bpe_ranks.keys()]
        lowercase__:       List[Any]        =       tokenizer.get_vocab()
        return cls(lowerCAmelCase__    ,			lowerCAmelCase__    ,			*lowerCAmelCase__    ,			**lowerCAmelCase__			)
 @classmethod
 def 		SCREAMING_SNAKE_CASE__       (       cls    ,			lowerCAmelCase__    ,			*lowerCAmelCase__    ,			**lowerCAmelCase__			) ->    List[Any]:




        '''simple docstring'''
        lowercase__:       int        =       GPTaTokenizer.from_pretrained(lowerCAmelCase__    ,			*lowerCAmelCase__    ,			**lowerCAmelCase__			)
        return cls.from_tokenizer(lowerCAmelCase__    ,			*lowerCAmelCase__    ,			**lowerCAmelCase__			)
 @classmethod
 def 		SCREAMING_SNAKE_CASE__       (       cls    ,			lowerCAmelCase__			) ->    Dict:




        '''simple docstring'''
        return cls(**lowerCAmelCase__			)
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    List[Any]:




        '''simple docstring'''
        return {
            "vocab": self.vocab,
            "merges": self.merges,
            "max_length": self.max_length,
            "pad_token_id": self.pad_token_id,
        }
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None			) ->    Optional[Any]:




        '''simple docstring'''
        lowercase__:       Optional[Any]        =       self.tf_tokenizer(lowerCAmelCase__			)
        lowercase__:       List[Any]        =       tf.ones_like(lowerCAmelCase__			)

        if self.pad_token_id is not None:
               # pad the tokens up to max length
               lowercase__:       int        =       max_length if max_length is not None else self.max_length

               if max_length is not None:
                      lowercase__      ,						lowercase__:       List[Any]        =       pad_model_inputs(
                          lowerCAmelCase__    ,			max_seq_length=lowerCAmelCase__    ,			pad_value=self.pad_token_id			)

        return {"attention_mask": attention_mask, "input_ids": input_ids}


