
'''simple docstring'''


import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   from transformers.models.bert.modeling_flax_bert import (
       FlaxBertForMaskedLM,
       FlaxBertForMultipleChoice,
       FlaxBertForNextSentencePrediction,
       FlaxBertForPreTraining,
       FlaxBertForQuestionAnswering,
       FlaxBertForSequenceClassification,
       FlaxBertForTokenClassification,
       FlaxBertModel,
   )


class 						lowercase__	(  unittest.TestCase   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__=1_3 ,   lowerCamelCase__=7 ,   lowerCamelCase__=True ,   lowerCamelCase__=True ,   lowerCamelCase__=True ,   lowerCamelCase__=True ,   lowerCamelCase__=9_9 ,   lowerCamelCase__=3_2 ,   lowerCamelCase__=5 ,   lowerCamelCase__=4 ,   lowerCamelCase__=3_7 ,   lowerCamelCase__="gelu" ,   lowerCamelCase__=0.1 ,   lowerCamelCase__=0.1 ,   lowerCamelCase__=5_1_2 ,   lowerCamelCase__=1_6 ,   lowerCamelCase__=2 ,   lowerCamelCase__=0.02 ,   lowerCamelCase__=4 ,   ):

      '''simple docstring'''







      UpperCamelCase       =  parent
      UpperCamelCase       =  batch_size
      UpperCamelCase       =  seq_length
      UpperCamelCase       =  is_training
      UpperCamelCase       =  use_attention_mask
      UpperCamelCase       =  use_token_type_ids
      UpperCamelCase       =  use_labels
      UpperCamelCase       =  vocab_size
      UpperCamelCase       =  hidden_size
      UpperCamelCase       =  num_hidden_layers
      UpperCamelCase       =  num_attention_heads
      UpperCamelCase       =  intermediate_size
      UpperCamelCase       =  hidden_act
      UpperCamelCase       =  hidden_dropout_prob
      UpperCamelCase       =  attention_probs_dropout_prob
      UpperCamelCase       =  max_position_embeddings
      UpperCamelCase       =  type_vocab_size
      UpperCamelCase       =  type_sequence_label_size
      UpperCamelCase       =  initializer_range
      UpperCamelCase       =  num_choices







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  ids_tensor([self.batch_size, self.seq_length] ,   self.vocab_size   )

      UpperCamelCase       =  None
      if self.use_attention_mask:
         UpperCamelCase       =  random_attention_mask([self.batch_size, self.seq_length]   )

      UpperCamelCase       =  None
      if self.use_token_type_ids:
         UpperCamelCase       =  ids_tensor([self.batch_size, self.seq_length] ,   self.type_vocab_size   )

      UpperCamelCase       =  BertConfig(
          vocab_size=self.vocab_size ,   hidden_size=self.hidden_size ,   num_hidden_layers=self.num_hidden_layers ,   num_attention_heads=self.num_attention_heads ,   intermediate_size=self.intermediate_size ,   hidden_act=self.hidden_act ,   hidden_dropout_prob=self.hidden_dropout_prob ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob ,   max_position_embeddings=self.max_position_embeddings ,   type_vocab_size=self.type_vocab_size ,   is_decoder=lowerCamelCase__ ,   initializer_range=self.initializer_range ,   )

      return config, input_ids, token_type_ids, attention_mask







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.prepare_config_and_inputs()
      UpperCamelCase     , UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  config_and_inputs
      UpperCamelCase       =  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
      return config, inputs_dict







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.prepare_config_and_inputs()
      UpperCamelCase     , UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  config_and_inputs

      UpperCamelCase       =  True
      UpperCamelCase       =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
      UpperCamelCase       =  ids_tensor([self.batch_size, self.seq_length] ,   vocab_size=2   )

      return (
          config,
          input_ids,
          attention_mask,
          encoder_hidden_states,
          encoder_attention_mask,
      )




@require_flax
class 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   _snake_case		   =							True

   _snake_case		   =							(
       (
           FlaxBertModel,
           FlaxBertForPreTraining,
           FlaxBertForMaskedLM,
           FlaxBertForMultipleChoice,
           FlaxBertForQuestionAnswering,
           FlaxBertForNextSentencePrediction,
           FlaxBertForSequenceClassification,
           FlaxBertForTokenClassification,
           FlaxBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  FlaxBertModelTester(self   )







   @slow
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  FlaxBertModel.from_pretrained('''bert-base-cased'''   )
      UpperCamelCase       =  model(np.ones((1, 1)   )   )
      self.assertIsNotNone(lowerCamelCase__   )







'''simple docstring'''


import os
import re
import unicodedata
from shutil import copyfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import is_torch_available, logging


if is_torch_available():
   import torch


if TYPE_CHECKING:
   from transformers.pipelines.conversational import Conversation


snake_case_ :  Union[str, Any]     =					logging.get_logger(__name__)
snake_case_ :  Union[str, Any]     =					{'vocab_file': 'spiece.model'}

snake_case_ :  Any     =					{
    'vocab_file': {
        'AI-Sweden/gpt-sw3-126m': 'https://huggingface.co/AI-Sweden/gpt-sw3-126m/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-350m': 'https://huggingface.co/AI-Sweden/gpt-sw3-350m/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-1.6b': 'https://huggingface.co/AI-Sweden/gpt-sw3-1.6b/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-6.7b': 'https://huggingface.co/AI-Sweden/gpt-sw3-6.7b/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-20b': 'https://huggingface.co/AI-Sweden/gpt-sw3-20b/resolve/main/spiece.model',
    }
}

snake_case_ :  Dict     =					{
    'AI-Sweden/gpt-sw3-126m': 2_048,
    'AI-Sweden/gpt-sw3-350m': 2_048,
    'AI-Sweden/gpt-sw3-1.6b': 2_048,
    'AI-Sweden/gpt-sw3-6.7b': 2_048,
    'AI-Sweden/gpt-sw3-20b': 2_048,
}


class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   _snake_case		   =							VOCAB_FILES_NAMES
   _snake_case		   =							PRETRAINED_VOCAB_FILES_MAP
   _snake_case		   =							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   _snake_case		   =							['''input_ids''', '''attention_mask''']







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__=False ,   lowerCamelCase__=False ,   lowerCamelCase__=False ,   lowerCamelCase__=None ,   lowerCamelCase__=None ,   lowerCamelCase__=None ,   lowerCamelCase__=None ,   lowerCamelCase__ = None ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase       =  {} if sp_model_kwargs is None else sp_model_kwargs

      UpperCamelCase       =  kwargs.get('''name_or_path'''   )
      if name_or_path is None:
         logger.warning(
             '''name_or_path not provided, will work for all GPTSw3 models except gpt-sw3-7b,'''
             ''' you are testing the model, this can safely be ignored'''   )
         UpperCamelCase       =  '''None'''

      # Default definitions for our 2 tokenizer versions, with None-checks to enable proper testing
      UpperCamelCase       =  '''<|endoftext|>''' if eos_token is None else eos_token
      UpperCamelCase       =  '''<unk>''' if unk_token is None else unk_token
      if "gpt-sw3-7b" in name_or_path:
         UpperCamelCase       =  unk_token if pad_token is None else pad_token
         UpperCamelCase       =  eos_token if bos_token is None else bos_token
      else:
         UpperCamelCase       =  '''<pad>''' if pad_token is None else pad_token
         UpperCamelCase       =  '''<s>''' if bos_token is None else bos_token

      super().__init__(
          do_lower_case=lowerCamelCase__ ,   remove_space=lowerCamelCase__ ,   keep_accents=lowerCamelCase__ ,   bos_token=lowerCamelCase__ ,   eos_token=lowerCamelCase__ ,   unk_token=lowerCamelCase__ ,   pad_token=lowerCamelCase__ ,   sp_model_kwargs=self.sp_model_kwargs ,   **lowerCamelCase__ ,   )

      UpperCamelCase       =  do_lower_case
      UpperCamelCase       =  remove_space
      UpperCamelCase       =  keep_accents
      UpperCamelCase       =  vocab_file

      UpperCamelCase       =  spm.SentencePieceProcessor(**self.sp_model_kwargs   )
      self.sp_model.Load(lowerCamelCase__   )

      # Used for whitespace normalization in input texts
      # fmt : off
      UpperCamelCase       =  {''' ''', ''' ''', ''' ''', ''' ''', ''' ''', '''　''', ''' ''', ''' ''', ''' ''', ''' ''', '''￼''', ''''''}
      # fmt : on

      # Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
      UpperCamelCase       =  re.compile(
          f'[{"".join(map(lowerCamelCase__ ,   list(range(0 ,   9   )   ) + list(range(1_1 ,   3_2   )   ) + list(range(1_2_7 ,   1_6_0   )   ) + [1_6_0, 1_7_3, 8_2_0_3]   )   )}]'   )







   def __getstate__(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.__dict__.copy()
      UpperCamelCase       =  None
      return state







   def __setstate__(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  d

      # for backward compatibility
      if not hasattr(self ,   '''sp_model_kwargs'''   ):
         UpperCamelCase       =  {}

      UpperCamelCase       =  spm.SentencePieceProcessor(**self.sp_model_kwargs   )
      self.sp_model.Load(self.vocab_file   )







   @property
   # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.vocab_size
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return len(self.sp_model   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  self.non_printing_characters_re.sub('''''' ,   lowerCamelCase__   )

      # Normalize whitespaces
      UpperCamelCase       =  ''''''.join([char if char not in self.whitespaces else ''' ''' for char in text]   )

      # NFC Unicode normalization
      UpperCamelCase       =  unicodedata.normalize('''NFC''' ,   lowerCamelCase__   )
      return text







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  self.preprocess_text(lowerCamelCase__   )
      return self.sp_model.encode(lowerCamelCase__ ,   out_type=lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return self.sp_model.PieceToId(lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return self.sp_model.IdToPiece(lowerCamelCase__   )







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__   ):

      '''simple docstring'''







      return out_string







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  []
      UpperCamelCase       =  ''''''
      UpperCamelCase       =  False
      for token in tokens:
         # make sure that special tokens are not decoded using sentencepiece model
         if token in self.all_special_tokens:
            # TODO: Check if this is needed, as it ensures that decode(encode(doc)) != doc by adding extra whitespace in the decoded document
            if not prev_is_special:
               out_string += " "

            out_string += self.sp_model.decode(lowerCamelCase__   ) + token
            UpperCamelCase       =  True
            UpperCamelCase       =  []
         else:
            current_sub_tokens.append(lowerCamelCase__   )
            UpperCamelCase       =  False
      out_string += self.sp_model.decode(lowerCamelCase__   )

      return out_string







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  {self.convert_ids_to_tokens(lowerCamelCase__   ): i for i in range(self.vocab_size   )}
      vocab.update(self.added_tokens_encoder   )
      return vocab







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = None   ):

      '''simple docstring'''







      if not os.path.isdir(lowerCamelCase__   ):
         logger.error(f'Vocabulary path ({save_directory}) should be a directory'   )
         return
      UpperCamelCase       =  os.path.join(
          lowerCamelCase__ ,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']   )

      if os.path.abspath(self.vocab_file   ) != os.path.abspath(lowerCamelCase__   ) and os.path.isfile(self.vocab_file   ):
         copyfile(self.vocab_file ,   lowerCamelCase__   )
      elif not os.path.isfile(self.vocab_file   ):
         with open(lowerCamelCase__ ,   '''wb'''   ) as fi:
            UpperCamelCase       =  self.sp_model.serialized_model_proto()
            fi.write(lowerCamelCase__   )

      return (out_vocab_file,)







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = False   ):

      '''simple docstring'''







      if isinstance(lowerCamelCase__ ,   lowerCamelCase__   ):
         UpperCamelCase       =  self.preprocess_text(lowerCamelCase__   )
         UpperCamelCase       =  self.sp_model.encode(lowerCamelCase__   )
      else:
         UpperCamelCase       =  [self.preprocess_text(lowerCamelCase__   ) for t in text]
         UpperCamelCase       =  self.sp_model.encode(lowerCamelCase__   )

      if return_tensors is True or return_tensors == "pt":
         UpperCamelCase       =  torch.tensor(lowerCamelCase__   )

      return token_ids







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return self.sp_model.decode(lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  [f'User: {text}' if is_user else f'Bot: {text}' for is_user, text in conversation.iter_texts()]
      UpperCamelCase       =  (
          f'{self.eos_token}{self.bos_token}' + f'{self.bos_token}'.join(lowerCamelCase__   ) + f'{self.bos_token}Bot:'
      )
      return self.encode(text=lowerCamelCase__   )





