






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  (    A__       ,   A__       ,   A__       ):
 _lowerCAmelCase       :       List[Any]        = [r'''h\.\d+\.attn\.bias''', r'''h\.\d+\.attn\.masked_bias''']







 @register_to_config
 def __init__(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   Optional[int] = None							,   SCREAMING_SNAKE_CASE					:   int = 5_02_57							,   SCREAMING_SNAKE_CASE					:   int = 10_24							,   SCREAMING_SNAKE_CASE					:   int = 7_68							,   SCREAMING_SNAKE_CASE					:   int = 12							,   SCREAMING_SNAKE_CASE					:   int = 12							,   SCREAMING_SNAKE_CASE					:   Optional[int] = None							,   SCREAMING_SNAKE_CASE					:   str = "gelu_new"							,   SCREAMING_SNAKE_CASE					:   float = 0.1							,   SCREAMING_SNAKE_CASE					:   float = 0.1							,   SCREAMING_SNAKE_CASE					:   float = 0.1							,   SCREAMING_SNAKE_CASE					:   float = 1e-5							,   SCREAMING_SNAKE_CASE					:   float = 0.0_2							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   bool = False							,   SCREAMING_SNAKE_CASE					:   bool = False							,   ):


      '''simple docstring'''






      super().__init__()

      UpperCamelCase__			:   List[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.'						)

      UpperCamelCase__			:   List[Any]       				=  prefix_inner_dim
      UpperCamelCase__			:   Optional[int]       				=  prefix_hidden_dim

      UpperCamelCase__			:   int       				=  (
          nn.Linear(self.prefix_inner_dim							,   self.prefix_hidden_dim						)
          if self.prefix_hidden_dim is not None
          else nn.Identity()
      )
      UpperCamelCase__			:   Any       				=  (
          nn.Linear(self.prefix_hidden_dim							,   SCREAMING_SNAKE_CASE						) if self.prefix_hidden_dim is not None else nn.Identity()
      )

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







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   torch.Tensor							,   SCREAMING_SNAKE_CASE					:   torch.Tensor							,   SCREAMING_SNAKE_CASE					:   Optional[torch.Tensor] = None							,   SCREAMING_SNAKE_CASE					:   Optional[torch.Tensor] = None							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  self.transformer.transformer.wte(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  self.encode_prefix(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Dict       				=  self.decode_prefix(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   str       				=  torch.cat((prefix_embeds, embedding_text)							,   dim=1						)

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







 def 					__lowercase					(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   torch.device						):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   Dict						):


      '''simple docstring'''






      return self.encode_prefix(SCREAMING_SNAKE_CASE						)







 @torch.no_grad()
 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






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







 @torch.no_grad()
 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   List[Any]=None							,   SCREAMING_SNAKE_CASE					:   List[str]=None							,   SCREAMING_SNAKE_CASE					:   str=None							,   SCREAMING_SNAKE_CASE					:   int = 5							,   SCREAMING_SNAKE_CASE					:   int = 67							,   SCREAMING_SNAKE_CASE					:   float = 1.0							,   SCREAMING_SNAKE_CASE					:   Optional[int] = None							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  eos_token_id
      UpperCamelCase__			:   Union[str, Any]       				=  None
      UpperCamelCase__			:   int       				=  None
      UpperCamelCase__			:   Tuple       				=  torch.ones(SCREAMING_SNAKE_CASE							,   device=SCREAMING_SNAKE_CASE							,   dtype=torch.int						)
      UpperCamelCase__			:   Dict       				=  torch.zeros(SCREAMING_SNAKE_CASE							,   device=SCREAMING_SNAKE_CASE							,   dtype=torch.bool						)

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

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

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

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

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







import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_blenderbot import BlenderbotTokenizer


if TYPE_CHECKING:
 from transformers.pipelines.conversational import Conversation

lowerCamelCase	:							Union[str, Any]				=logging.get_logger(__name__)


lowerCamelCase	:							str				={
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

lowerCamelCase	:							str				={
    '''vocab_file''': {'''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/vocab.json'''},
    '''merges_file''': {'''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/merges.txt'''},
    '''tokenizer_config_file''': {
        '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/tokenizer_config.json'''
    },
}

lowerCamelCase	:							str				={'''facebook/blenderbot-3B''': 128}





class 			__a  (    A__       ):
 _lowerCAmelCase       :       Union[str, Any]        = VOCAB_FILES_NAMES
 _lowerCAmelCase       :       Dict        = PRETRAINED_VOCAB_FILES_MAP
 _lowerCAmelCase       :       List[Any]        = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 _lowerCAmelCase       :       int        = ['''input_ids''', '''attention_mask''']
 _lowerCAmelCase       :       Tuple        = BlenderbotTokenizer







 def __init__(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Optional[int]=None							,   SCREAMING_SNAKE_CASE					:   Dict=None							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]=None							,   SCREAMING_SNAKE_CASE					:   Tuple="replace"							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]="<s>"							,   SCREAMING_SNAKE_CASE					:   str="</s>"							,   SCREAMING_SNAKE_CASE					:   str="</s>"							,   SCREAMING_SNAKE_CASE					:   Tuple="<s>"							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]="<unk>"							,   SCREAMING_SNAKE_CASE					:   str="<pad>"							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]="<mask>"							,   SCREAMING_SNAKE_CASE					:   Any=False							,   SCREAMING_SNAKE_CASE					:   List[Any]=True							,   **SCREAMING_SNAKE_CASE					:   int							,   ):


      '''simple docstring'''






      super().__init__(
          SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   tokenizer_file=SCREAMING_SNAKE_CASE							,   errors=SCREAMING_SNAKE_CASE							,   bos_token=SCREAMING_SNAKE_CASE							,   eos_token=SCREAMING_SNAKE_CASE							,   sep_token=SCREAMING_SNAKE_CASE							,   cls_token=SCREAMING_SNAKE_CASE							,   unk_token=SCREAMING_SNAKE_CASE							,   pad_token=SCREAMING_SNAKE_CASE							,   mask_token=SCREAMING_SNAKE_CASE							,   add_prefix_space=SCREAMING_SNAKE_CASE							,   trim_offsets=SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE							,   )

      UpperCamelCase__			:   List[Any]       				=  json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()						)
      if pre_tok_state.get("add_prefix_space"							,   SCREAMING_SNAKE_CASE						) != add_prefix_space:
           UpperCamelCase__			:   str       				=  getattr(SCREAMING_SNAKE_CASE							,   pre_tok_state.pop("type"						)						)
           UpperCamelCase__			:   Union[str, Any]       				=  add_prefix_space
           UpperCamelCase__			:   Union[str, Any]       				=  pre_tok_class(**SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Tuple       				=  add_prefix_space

      UpperCamelCase__			:   Optional[int]       				=  "post_processor"
      UpperCamelCase__			:   Union[str, Any]       				=  getattr(self.backend_tokenizer							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      if tokenizer_component_instance:
           UpperCamelCase__			:   Tuple       				=  json.loads(tokenizer_component_instance.__getstate__()						)

           # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
           if "sep" in state:
                UpperCamelCase__			:   Tuple       				=  tuple(state["sep"]						)
           if "cls" in state:
                UpperCamelCase__			:   Optional[int]       				=  tuple(state["cls"]						)

           UpperCamelCase__			:   List[Any]       				=  False

           if state.get("add_prefix_space"							,   SCREAMING_SNAKE_CASE						) != add_prefix_space:
                UpperCamelCase__			:   str       				=  add_prefix_space
                UpperCamelCase__			:   Optional[Any]       				=  True

           if state.get("trim_offsets"							,   SCREAMING_SNAKE_CASE						) != trim_offsets:
                UpperCamelCase__			:   Optional[Any]       				=  trim_offsets
                UpperCamelCase__			:   Tuple       				=  True

           if changes_to_apply:
                UpperCamelCase__			:   Optional[int]       				=  getattr(SCREAMING_SNAKE_CASE							,   state.pop("type"						)						)
                UpperCamelCase__			:   Dict       				=  component_class(**SCREAMING_SNAKE_CASE						)
                setattr(self.backend_tokenizer							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)







 @property
 # Copied from transformers.models.roberta.tokenization_roberta_fast.RobertaTokenizerFast.mask_token with Roberta->Blenderbot, RoBERTa->Blenderbot
 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      if self._mask_token is None:
           if self.verbose:
                logger.error("Using mask_token, but it is not set yet."						)
           return None
      return str(self._mask_token						)







 @mask_token.setter
 def 					__lowercase					(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  AddedToken(SCREAMING_SNAKE_CASE							,   lstrip=SCREAMING_SNAKE_CASE							,   rstrip=SCREAMING_SNAKE_CASE						) if isinstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						) else value
      UpperCamelCase__			:   str       				=  value







 def 					__lowercase					(		self					:   Any							,   *SCREAMING_SNAKE_CASE					:   List[str]							,   **SCREAMING_SNAKE_CASE					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  kwargs.get("is_split_into_words"							,   SCREAMING_SNAKE_CASE						)
      assert self.add_prefix_space or not is_split_into_words, (
          F'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
          "to use it with pretokenized inputs."
      )

      return super()._batch_encode_plus(*SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Dict							,   *SCREAMING_SNAKE_CASE					:   int							,   **SCREAMING_SNAKE_CASE					:   Dict						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  kwargs.get("is_split_into_words"							,   SCREAMING_SNAKE_CASE						)

      assert self.add_prefix_space or not is_split_into_words, (
          F'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
          "to use it with pretokenized inputs."
      )

      return super()._encode_plus(*SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   Optional[str] = None						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[Any]       				=  self._tokenizer.model.save(SCREAMING_SNAKE_CASE							,   name=SCREAMING_SNAKE_CASE						)
      return tuple(SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   List[int]							,   SCREAMING_SNAKE_CASE					:   Optional[List[int]] = None						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[Any]       				=  [self.sep_token_id]
      UpperCamelCase__			:   int       				=  [self.cls_token_id]

      if token_ids_a is None:
           return len(cls + token_ids_a + sep						) * [0]
      return len(cls + token_ids_a + sep + sep + token_ids_a + sep						) * [0]







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   List[int]							,   SCREAMING_SNAKE_CASE					:   Optional[List[int]] = None						):


      '''simple docstring'''






      return token_ids_a + [self.eos_token_id]







 def 					__lowercase					(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   "Conversation"						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[Any]       				=  []
      for is_user, text in conversation.iter_texts():
           if is_user:
                # We need to space prefix as it's being done within blenderbot
                inputs.append(" " + text						)
           else:
                # Generated responses should contain them already.
                inputs.append(SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   List[str]       				=  "  ".join(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Optional[Any]       				=  self.encode(SCREAMING_SNAKE_CASE						)
      if len(SCREAMING_SNAKE_CASE						) > self.model_max_length:
           UpperCamelCase__			:   Dict       				=  input_ids[-self.model_max_length :]
           logger.warning(F'Trimmed input from conversation as it was longer than {self.model_max_length} tokens.'						)
      return input_ids