'''simple docstring'''

import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      FalconForCausalLM,
      FalconForQuestionAnswering,
      FalconForSequenceClassification,
      FalconForTokenClassification,
      FalconModel,
  )


class 				_lowercase	:
      '''simple docstring'''


      def __init__(				self					,lowerCamelCase_					,lowerCamelCase_=3					,lowerCamelCase_=7					,lowerCamelCase_=True					,lowerCamelCase_=True					,lowerCamelCase_=False					,lowerCamelCase_=True					,lowerCamelCase_=99					,lowerCamelCase_=32					,lowerCamelCase_=5					,lowerCamelCase_=4					,lowerCamelCase_=37					,lowerCamelCase_="gelu"					,lowerCamelCase_=0.1					,lowerCamelCase_=0.1					,lowerCamelCase_=512					,lowerCamelCase_=16					,lowerCamelCase_=2					,lowerCamelCase_=0.02					,lowerCamelCase_=3					,lowerCamelCase_=4					,lowerCamelCase_=None					,)			->			int:



         '''simple docstring'''



         UpperCAmelCase__							:				Dict   				=  parent
         UpperCAmelCase__							:				List[Any]   				=  batch_size
         UpperCAmelCase__							:				List[Any]   				=  seq_length
         UpperCAmelCase__							:				List[str]   				=  is_training
         UpperCAmelCase__							:				Tuple   				=  use_input_mask
         UpperCAmelCase__							:				Dict   				=  use_token_type_ids
         UpperCAmelCase__							:				Tuple   				=  use_labels
         UpperCAmelCase__							:				Optional[Any]   				=  vocab_size
         UpperCAmelCase__							:				Union[str, Any]   				=  hidden_size
         UpperCAmelCase__							:				List[str]   				=  num_hidden_layers
         UpperCAmelCase__							:				List[Any]   				=  num_attention_heads
         UpperCAmelCase__							:				Dict   				=  intermediate_size
         UpperCAmelCase__							:				Tuple   				=  hidden_act
         UpperCAmelCase__							:				Optional[int]   				=  hidden_dropout_prob
         UpperCAmelCase__							:				Optional[Any]   				=  attention_probs_dropout_prob
         UpperCAmelCase__							:				List[str]   				=  max_position_embeddings
         UpperCAmelCase__							:				List[str]   				=  type_vocab_size
         UpperCAmelCase__							:				Tuple   				=  type_sequence_label_size
         UpperCAmelCase__							:				Optional[Any]   				=  initializer_range
         UpperCAmelCase__							:				str   				=  num_labels
         UpperCAmelCase__							:				List[Any]   				=  num_choices
         UpperCAmelCase__							:				Any   				=  scope




      def 		lowerCAmelCase__					(				self							)			->			Any:



         '''simple docstring'''



         UpperCAmelCase__							:				int   				=  ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size							)

         UpperCAmelCase__							:				Dict   				=  None
         if self.use_input_mask:
            UpperCAmelCase__							:				Union[str, Any]   				=  random_attention_mask([self.batch_size, self.seq_length]							)

         UpperCAmelCase__							:				Dict   				=  None

         UpperCAmelCase__							:				Union[str, Any]   				=  None
         UpperCAmelCase__							:				Optional[Any]   				=  None
         UpperCAmelCase__							:				Optional[int]   				=  None
         if self.use_labels:
            UpperCAmelCase__							:				Tuple   				=  ids_tensor([self.batch_size]					,self.type_sequence_label_size							)
            UpperCAmelCase__							:				Optional[Any]   				=  ids_tensor([self.batch_size, self.seq_length]					,self.num_labels							)
            UpperCAmelCase__							:				List[Any]   				=  ids_tensor([self.batch_size]					,self.num_choices							)

         UpperCAmelCase__							:				List[str]   				=  self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




      def 		lowerCAmelCase__					(				self							)			->			Optional[Any]:



         '''simple docstring'''



         return FalconConfig(
             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					,pad_token_id=1					,new_decoder_architecture=lowerCamelCase_					,)




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_							)			->			Any:



         '''simple docstring'''



         UpperCAmelCase__							:				Optional[int]   				=  FalconModel(config=lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()
         UpperCAmelCase__							:				List[str]   				=  model(lowerCamelCase_					,attention_mask=lowerCamelCase_							)
         UpperCAmelCase__							:				Any   				=  model(lowerCamelCase_							)
         self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)							)




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,)			->			Dict:



         '''simple docstring'''



         UpperCAmelCase__							:				Optional[Any]   				=  True
         UpperCAmelCase__							:				List[str]   				=  FalconModel(lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()
         UpperCAmelCase__							:				List[str]   				=  model(
             lowerCamelCase_					,attention_mask=lowerCamelCase_					,encoder_hidden_states=lowerCamelCase_					,encoder_attention_mask=lowerCamelCase_					,)
         UpperCAmelCase__							:				List[Any]   				=  model(
             lowerCamelCase_					,attention_mask=lowerCamelCase_					,encoder_hidden_states=lowerCamelCase_					,)
         UpperCAmelCase__							:				List[Any]   				=  model(lowerCamelCase_					,attention_mask=lowerCamelCase_							)
         self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)							)




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,)			->			int:



         '''simple docstring'''



         UpperCAmelCase__							:				Dict   				=  FalconForCausalLM(config=lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()
         UpperCAmelCase__							:				Union[str, Any]   				=  model(lowerCamelCase_					,attention_mask=lowerCamelCase_					,labels=lowerCamelCase_							)
         self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.vocab_size)							)




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,)			->			Optional[int]:



         '''simple docstring'''



         UpperCAmelCase__							:				Optional[Any]   				=  True
         UpperCAmelCase__							:				str   				=  True
         UpperCAmelCase__							:				Dict   				=  FalconForCausalLM(config=lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()

         # first forward pass
         UpperCAmelCase__							:				Optional[Any]   				=  model(
             lowerCamelCase_					,attention_mask=lowerCamelCase_					,encoder_hidden_states=lowerCamelCase_					,encoder_attention_mask=lowerCamelCase_					,use_cache=lowerCamelCase_					,)
         UpperCAmelCase__							:				Optional[Any]   				=  outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         UpperCAmelCase__							:				str   				=  ids_tensor((self.batch_size, 3)					,config.vocab_size							)
         UpperCAmelCase__							:				Optional[Any]   				=  ids_tensor((self.batch_size, 3)					,vocab_size=2							)

         # append to next input_ids and
         UpperCAmelCase__							:				int   				=  torch.cat([input_ids, next_tokens]					,dim=-1							)
         UpperCAmelCase__							:				Union[str, Any]   				=  torch.cat([input_mask, next_mask]					,dim=-1							)

         UpperCAmelCase__							:				Optional[Any]   				=  model(
             lowerCamelCase_					,attention_mask=lowerCamelCase_					,encoder_hidden_states=lowerCamelCase_					,encoder_attention_mask=lowerCamelCase_					,output_hidden_states=lowerCamelCase_					,)['''hidden_states'''][0]
         UpperCAmelCase__							:				Union[str, Any]   				=  model(
             lowerCamelCase_					,attention_mask=lowerCamelCase_					,encoder_hidden_states=lowerCamelCase_					,encoder_attention_mask=lowerCamelCase_					,past_key_values=lowerCamelCase_					,output_hidden_states=lowerCamelCase_					,)['''hidden_states'''][0]

         # select random slice
         UpperCAmelCase__							:				Any   				=  ids_tensor((1,)					,output_from_past.shape[-1]							).item()
         UpperCAmelCase__							:				Union[str, Any]   				=  output_from_no_past[:, -3:, random_slice_idx].detach()
         UpperCAmelCase__							:				List[Any]   				=  output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]							)

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(lowerCamelCase_					,lowerCamelCase_					,atol=1e-3							)							)




      def 		lowerCAmelCase__					(				self							)			->			Dict:



         '''simple docstring'''



         UpperCAmelCase__							:				List[str]   				=  self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase__
         )				,      (
             UpperCAmelCase__
         )				,      (
             UpperCAmelCase__
         )				,      (
             UpperCAmelCase__
         )				,      (
             UpperCAmelCase__
         )				,      (
             UpperCAmelCase__
         )				,      (
             UpperCAmelCase__
         )				,      
         )							:				Dict   				=  config_and_inputs
         UpperCAmelCase__							:				List[Any]   				=  {'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict




@require_torch
class 				_lowercase	(      lowerCAmelCase		,lowerCAmelCase		,lowerCAmelCase		,unittest.TestCase   ):
      '''simple docstring'''


      UpperCAmelCase_      :      List[Any]      							=   (
          (
              FalconModel,
              FalconForCausalLM,
              FalconForSequenceClassification,
              FalconForTokenClassification,
              FalconForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      UpperCAmelCase_      :      List[str]      							=   (FalconForCausalLM,) if is_torch_available() else ()
      UpperCAmelCase_      :      Optional[Any]      							=   (
          {
              '''feature-extraction''': FalconModel,
              '''text-classification''': FalconForSequenceClassification,
              '''text-generation''': FalconForCausalLM,
              '''question-answering''': FalconForQuestionAnswering,
              '''token-classification''': FalconForTokenClassification,
              '''zero-shot''': FalconForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      UpperCAmelCase_      :      List[str]      							=   False
      UpperCAmelCase_      :      Optional[Any]      							=   False




      def 		lowerCAmelCase__					(				self							)			->			str:



         '''simple docstring'''



         UpperCAmelCase__							:				Any   				=  FalconModelTester(self							)
         UpperCAmelCase__							:				Any   				=  ConfigTester(self					,config_class=lowerCamelCase_					,hidden_size=37							)




      def 		lowerCAmelCase__					(				self							)			->			Tuple:



         '''simple docstring'''



         self.config_tester.run_common_tests()




      def 		lowerCAmelCase__					(				self							)			->			List[Any]:



         '''simple docstring'''



         UpperCAmelCase__							:				Optional[Any]   				=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowerCamelCase_							)




      def 		lowerCAmelCase__					(				self							)			->			List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				,      *UpperCAmelCase__							:				Optional[int]   				=  self.model_tester.prepare_config_and_inputs()
         for alibi in [True, False]:
            UpperCAmelCase__							:				Optional[Any]   				=  alibi
            self.model_tester.create_and_check_model(lowerCamelCase_					,*lowerCamelCase_							)




      def 		lowerCAmelCase__					(				self							)			->			Tuple:



         '''simple docstring'''



         UpperCAmelCase__				,      UpperCAmelCase__							:				int   				=  self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase__							:				List[Any]   				=  3
         UpperCAmelCase__							:				Union[str, Any]   				=  input_dict['''input_ids''']
         UpperCAmelCase__							:				int   				=  input_ids.ne(1							).to(lowerCamelCase_							)
         UpperCAmelCase__							:				List[str]   				=  ids_tensor([self.model_tester.batch_size]					,self.model_tester.type_sequence_label_size							)
         UpperCAmelCase__							:				Any   				=  FalconForSequenceClassification(lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()
         UpperCAmelCase__							:				Any   				=  model(lowerCamelCase_					,attention_mask=lowerCamelCase_					,labels=lowerCamelCase_							)
         self.assertEqual(result.logits.shape					,(self.model_tester.batch_size, self.model_tester.num_labels)							)




      def 		lowerCAmelCase__					(				self							)			->			List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				,      UpperCAmelCase__							:				Optional[int]   				=  self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase__							:				Optional[int]   				=  3
         UpperCAmelCase__							:				Optional[int]   				=  '''single_label_classification'''
         UpperCAmelCase__							:				Tuple   				=  input_dict['''input_ids''']
         UpperCAmelCase__							:				Optional[int]   				=  input_ids.ne(1							).to(lowerCamelCase_							)
         UpperCAmelCase__							:				Union[str, Any]   				=  ids_tensor([self.model_tester.batch_size]					,self.model_tester.type_sequence_label_size							)
         UpperCAmelCase__							:				Union[str, Any]   				=  FalconForSequenceClassification(lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()
         UpperCAmelCase__							:				Any   				=  model(lowerCamelCase_					,attention_mask=lowerCamelCase_					,labels=lowerCamelCase_							)
         self.assertEqual(result.logits.shape					,(self.model_tester.batch_size, self.model_tester.num_labels)							)




      def 		lowerCAmelCase__					(				self							)			->			List[str]:



         '''simple docstring'''



         UpperCAmelCase__				,      UpperCAmelCase__							:				Optional[int]   				=  self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase__							:				Any   				=  input_dict['''input_ids''']
         UpperCAmelCase__							:				Optional[Any]   				=  FalconForCausalLM(lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()
         UpperCAmelCase__							:				Dict   				=  model(lowerCamelCase_					,use_cache=lowerCamelCase_							)
         UpperCAmelCase__							:				List[Any]   				=  input_ids.shape[0]
         UpperCAmelCase__							:				Optional[int]   				=  model._convert_to_rw_cache(result.past_key_values							)
         UpperCAmelCase__							:				Tuple   				=  model._convert_cache_to_standard_format(lowerCamelCase_					,lowerCamelCase_							)
         for layer in range(len(lowerCamelCase_							)							):
            for tensor_idx in range(2							):
               self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3							)
               self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4							)
               self.assertTrue(
                   torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]							)							)




      def 		lowerCAmelCase__					(				self							)			->			Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				,      UpperCAmelCase__							:				Optional[Any]   				=  self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase__							:				str   				=  3
         UpperCAmelCase__							:				List[Any]   				=  '''multi_label_classification'''
         UpperCAmelCase__							:				List[str]   				=  input_dict['''input_ids''']
         UpperCAmelCase__							:				Optional[Any]   				=  input_ids.ne(1							).to(lowerCamelCase_							)
         UpperCAmelCase__							:				Optional[Any]   				=  ids_tensor(
             [self.model_tester.batch_size, config.num_labels]					,self.model_tester.type_sequence_label_size							).to(torch.float							)
         UpperCAmelCase__							:				int   				=  FalconForSequenceClassification(lowerCamelCase_							)
         model.to(lowerCamelCase_							)
         model.eval()
         UpperCAmelCase__							:				List[Any]   				=  model(lowerCamelCase_					,attention_mask=lowerCamelCase_					,labels=lowerCamelCase_							)
         self.assertEqual(result.logits.shape					,(self.model_tester.batch_size, self.model_tester.num_labels)							)




      def 		lowerCAmelCase__					(				self							)			->			Optional[int]:



         '''simple docstring'''



         for model_class in self.all_generative_model_classes:
            UpperCAmelCase__				,      UpperCAmelCase__							:				Any   				=  self.model_tester.prepare_config_and_inputs_for_common()

            # If it doesn't support cache, pass the test
            if not hasattr(lowerCamelCase_					,'''use_cache'''							):
               return

            UpperCAmelCase__							:				Union[str, Any]   				=  model_class(lowerCamelCase_							).to(lowerCamelCase_							)
            if "use_cache" not in inputs:
               UpperCAmelCase__							:				Union[str, Any]   				=  True
            UpperCAmelCase__							:				Dict   				=  model(**lowerCamelCase_							)

            # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
            if "past_key_values" not in outputs:
               return

            UpperCAmelCase__							:				Any   				=  (
                getattr(lowerCamelCase_					,'''decoder_layers'''					,lowerCamelCase_							)
                or getattr(lowerCamelCase_					,'''num_decoder_layers'''					,lowerCamelCase_							)
                or config.num_hidden_layers
            )
            UpperCAmelCase__							:				Optional[Any]   				=  getattr(lowerCamelCase_					,'''num_kv_heads'''					,config.num_attention_heads							)
            UpperCAmelCase__							:				Tuple   				=  getattr(lowerCamelCase_					,'''d_model'''					,config.hidden_size							)
            UpperCAmelCase__							:				str   				=  embed_dim // num_attention_heads

            UpperCAmelCase__							:				str   				=  outputs['''past_key_values''']
            self.assertEqual(len(lowerCamelCase_							)					,lowerCamelCase_							)

            UpperCAmelCase__				,      UpperCAmelCase__							:				Optional[int]   				=  inputs['''input_ids'''].shape
            for i in range(lowerCamelCase_							):
               if config.new_decoder_architecture:
                  UpperCAmelCase__							:				Any   				=  config.num_attention_heads
               elif config.multi_query:
                  UpperCAmelCase__							:				int   				=  1
               self.assertEqual(len(past_kv[0]							)					,2							)  # K V for the decoder = 2
               self.assertEqual(
                   past_kv[i][0].shape					,(batch_size, num_attention_heads, seq_length, per_head_embed_dim)							)
               self.assertEqual(
                   past_kv[i][1].shape					,(batch_size, num_attention_heads, seq_length, per_head_embed_dim)							)




@require_torch
class 				_lowercase	(      unittest.TestCase   ):
      '''simple docstring'''


      @slow
      def 		lowerCAmelCase__					(				self							)			->			Optional[int]:



         '''simple docstring'''



         UpperCAmelCase__							:				Union[str, Any]   				=  AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b'''							)
         UpperCAmelCase__							:				Union[str, Any]   				=  FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b'''							)
         model.eval()
         model.to(lowerCamelCase_							)
         UpperCAmelCase__							:				str   				=  tokenizer('''My favorite food is'''					,return_tensors='''pt'''							).to(lowerCamelCase_							)

         UpperCAmelCase__							:				Union[str, Any]   				=  (
             '''My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'''
         )

         UpperCAmelCase__							:				Tuple   				=  model.generate(**lowerCamelCase_					,do_sample=lowerCamelCase_					,max_new_tokens=19							)
         UpperCAmelCase__							:				List[Any]   				=  tokenizer.batch_decode(lowerCamelCase_							)[0]

         self.assertEqual(lowerCamelCase_					,lowerCamelCase_							)




      @slow
      def 		lowerCAmelCase__					(				self							)			->			Optional[Any]:



         '''simple docstring'''



         for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
            UpperCAmelCase__							:				Optional[Any]   				=  AutoTokenizer.from_pretrained(lowerCamelCase_							)
            UpperCAmelCase__							:				List[str]   				=  FalconForCausalLM.from_pretrained(lowerCamelCase_							)
            model.eval()
            model.to(lowerCamelCase_							)
            UpperCAmelCase__							:				Optional[Any]   				=  tokenizer('''My favorite food is'''					,return_tensors='''pt'''							).to(lowerCamelCase_							)

            # We just test that these run without errors - the models are randomly initialized
            # and so the actual text outputs will be garbage
            model.generate(**lowerCamelCase_					,do_sample=lowerCamelCase_					,max_new_tokens=4							)
            model.generate(**lowerCamelCase_					,do_sample=lowerCamelCase_					,max_new_tokens=4							)
            model.generate(**lowerCamelCase_					,num_beams=2					,max_new_tokens=4							)




      @slow
      def 		lowerCAmelCase__					(				self							)			->			Tuple:



         '''simple docstring'''



         with torch.no_grad():
            for repo in [
                "Rocketknight1/falcon-rw-1b",
                "Rocketknight1/tiny-random-falcon-7b",
                "Rocketknight1/tiny-random-falcon-40b",
            ]:
               UpperCAmelCase__							:				int   				=  AutoTokenizer.from_pretrained(lowerCamelCase_							)
               UpperCAmelCase__							:				List[str]   				=  FalconForCausalLM.from_pretrained(lowerCamelCase_							)
               model.eval()
               model.to(device=lowerCamelCase_							)
               UpperCAmelCase__							:				str   				=  tokenizer('''My favorite food is'''					,return_tensors='''pt'''							).to(lowerCamelCase_							)

               # Test results are the same with and without cache
               UpperCAmelCase__							:				Optional[Any]   				=  model.generate(**lowerCamelCase_					,do_sample=lowerCamelCase_					,max_new_tokens=20					,use_cache=lowerCamelCase_							)
               UpperCAmelCase__							:				Any   				=  model.generate(**lowerCamelCase_					,do_sample=lowerCamelCase_					,max_new_tokens=20					,use_cache=lowerCamelCase_							)
               self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0							)





'''simple docstring'''

import os
from shutil import copyfile
from typing import List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


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

UpperCamelCase__     :	str				=			{'vocab_file': 'sentencepiece.model'}

UpperCamelCase__     :	int				=			{
    'vocab_file': {
        'google/rembert': 'https://huggingface.co/google/rembert/resolve/main/sentencepiece.model',
    },
}

UpperCamelCase__     :	Any				=			{
    'google/rembert': 256,
}


class 				_lowercase	(      lowerCAmelCase   ):
      '''simple docstring'''


      UpperCAmelCase_      :      List[Any]      							=   VOCAB_FILES_NAMES
      UpperCAmelCase_      :      List[Any]      							=   PRETRAINED_VOCAB_FILES_MAP
      UpperCAmelCase_      :      Tuple      							=   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES




      def __init__(				self					,lowerCamelCase_					,lowerCamelCase_=False					,lowerCamelCase_=True					,lowerCamelCase_=True					,lowerCamelCase_="[CLS]"					,lowerCamelCase_="[SEP]"					,lowerCamelCase_="[UNK]"					,lowerCamelCase_="[SEP]"					,lowerCamelCase_="[PAD]"					,lowerCamelCase_="[CLS]"					,lowerCamelCase_="[MASK]"					,**lowerCamelCase_					,)			->			Tuple:



         '''simple docstring'''



         super().__init__(
             do_lower_case=lowerCamelCase_					,remove_space=lowerCamelCase_					,keep_accents=lowerCamelCase_					,bos_token=lowerCamelCase_					,eos_token=lowerCamelCase_					,unk_token=lowerCamelCase_					,sep_token=lowerCamelCase_					,pad_token=lowerCamelCase_					,cls_token=lowerCamelCase_					,mask_token=lowerCamelCase_					,**lowerCamelCase_					,)

         UpperCAmelCase__							:				Union[str, Any]   				=  do_lower_case
         UpperCAmelCase__							:				str   				=  remove_space
         UpperCAmelCase__							:				Optional[Any]   				=  keep_accents
         UpperCAmelCase__							:				List[Any]   				=  vocab_file

         UpperCAmelCase__							:				Dict   				=  spm.SentencePieceProcessor()
         self.sp_model.Load(lowerCamelCase_							)




      @property
      def 		lowerCAmelCase__					(				self							)			->			Optional[Any]:



         '''simple docstring'''



         return len(self.sp_model							)




      def 		lowerCAmelCase__					(				self							)			->			Dict:



         '''simple docstring'''



         UpperCAmelCase__							:				Optional[int]   				=  {self.convert_ids_to_tokens(lowerCamelCase_							): i for i in range(self.vocab_size							)}
         vocab.update(self.added_tokens_encoder							)
         return vocab




      def __getstate__(				self							)			->			Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__							:				Optional[int]   				=  self.__dict__.copy()
         UpperCAmelCase__							:				Any   				=  None
         return state




      def __setstate__(				self					,lowerCamelCase_							)			->			Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__							:				List[str]   				=  d
         UpperCAmelCase__							:				Any   				=  spm.SentencePieceProcessor()
         self.sp_model.Load(self.vocab_file							)




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_=False							)			->			Tuple:



         '''simple docstring'''



         UpperCAmelCase__							:				str   				=  self.sp_model.EncodeAsPieces(lowerCamelCase_							)
         return pieces




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_							)			->			str:



         '''simple docstring'''



         return self.sp_model.PieceToId(lowerCamelCase_							)




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_							)			->			Optional[int]:



         '''simple docstring'''



         return self.sp_model.IdToPiece(lowerCamelCase_							)




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_							)			->			List[Any]:



         '''simple docstring'''



         UpperCAmelCase__							:				List[Any]   				=  self.sp_model.decode_pieces(lowerCamelCase_							)
         return out_string




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_ = None							)			->			List[int]:



         '''simple docstring'''



         UpperCAmelCase__							:				Union[str, Any]   				=  [self.sep_token_id]
         UpperCAmelCase__							:				List[Any]   				=  [self.cls_token_id]
         if token_ids_a is None:
            return cls + token_ids_a + sep
         return cls + token_ids_a + sep + token_ids_a + sep




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_ = None					,lowerCamelCase_ = False							)			->			List[int]:



         '''simple docstring'''



         if already_has_special_tokens:
            if token_ids_a is not None:
               raise ValueError(
                   '''You should not supply a second sequence if the provided sequence of '''
                   '''ids is already formatted with special tokens for the model.'''							)
            return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

         if token_ids_a is not None:
            return [1] + ([0] * len(lowerCamelCase_							)) + [1] + ([0] * len(lowerCamelCase_							)) + [1]
         return [1] + ([0] * len(lowerCamelCase_							)) + [1]




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_ = None							)			->			List[int]:



         '''simple docstring'''



         UpperCAmelCase__							:				Tuple   				=  [self.sep_token_id]
         UpperCAmelCase__							:				List[Any]   				=  [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							) * [0] + len(token_ids_a + sep							) * [1]




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_ = None							)			->			Tuple[str]:



         '''simple docstring'''



         if not os.path.isdir(lowerCamelCase_							):
            logger.error('''Vocabulary path ({}) should be a directory'''.format(lowerCamelCase_							)							)
            return
         UpperCAmelCase__							:				Union[str, Any]   				=  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_							):
            copyfile(self.vocab_file					,lowerCamelCase_							)

         return (out_vocab_file,)




