






"""simple docstring"""





import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
  import os

  import jax.numpy as jnp
  from jax import jit

  from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
  from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

  _SCREAMING_SNAKE_CASE					:				List[Any]					    =						'0.12'  # assumed parallelism: 8


if is_torch_available():
  import torch







def      lowerCamelCase__       (							_lowerCamelCase       :		Dict		,				_lowerCamelCase       :		Union[str, Any]		,				_lowerCamelCase       :		Optional[int]=None )      ->				Any:







  if rng is None:
    lowerCamelCase_         =						random.Random()

  lowerCamelCase_         =						1
  for dim in shape:
    total_dims *= dim

  lowerCamelCase_         =						[]
  for _ in range(_lowerCamelCase ):
    values.append(rng.randint(0		,				vocab_size - 1 ) )

  lowerCamelCase_         =						np.array(_lowerCamelCase		,				dtype=jnp.intaa ).reshape(_lowerCamelCase )

  return output







def      lowerCamelCase__       (							_lowerCamelCase       :		Dict		,				_lowerCamelCase       :		Union[str, Any]=None )      ->				Tuple:
  lowerCamelCase_         =						ids_tensor(_lowerCamelCase		,				vocab_size=2		,				rng=_lowerCamelCase )
  # make sure that at least one token is attended to for each batch
  lowerCamelCase_         =						1
  return attn_mask






@require_flax
class      a						:
    SCREAMING_SNAKE_CASE						:     List[str]						= None
    SCREAMING_SNAKE_CASE						:     Any						= ()






    def    UpperCamelCase       (       self     :							int					)     ->		List[str]:
      lowerCamelCase_     ,	lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs_for_common()

      # cut to half length & take max batch_size 3
      lowerCamelCase_         =						2
      lowerCamelCase_         =						inputs['input_ids'].shape[-1] // 2
      lowerCamelCase_         =						inputs['input_ids'][:max_batch_size, :sequence_length]

      lowerCamelCase_         =						jnp.ones_like(__lowerCamelCase					)
      lowerCamelCase_         =						attention_mask[:max_batch_size, :sequence_length]

      # generate max 5 tokens
      lowerCamelCase_         =						input_ids.shape[-1] + 5
      if config.eos_token_id is not None and config.pad_token_id is None:
        # hack to allow generate for models such as GPT2 as is done in `generate()`
        lowerCamelCase_         =						config.eos_token_id
      return config, input_ids, attention_mask, max_length






    @is_pt_flax_cross_test
    def    UpperCamelCase       (       self     :							Optional[Any]					)     ->		int:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						False
      lowerCamelCase_         =						max_length
      lowerCamelCase_         =						0

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model_class.__name__[4:]  # Skip the "Flax" at the beginning
        lowerCamelCase_         =						getattr(__lowerCamelCase       ,						__lowerCamelCase					)
        lowerCamelCase_         =						pt_model_class(__lowerCamelCase					).eval()
        lowerCamelCase_         =						load_flax_weights_in_pytorch_model(__lowerCamelCase       ,						flax_model.params					)

        lowerCamelCase_         =						flax_model.generate(__lowerCamelCase					).sequences
        lowerCamelCase_         =						pt_model.generate(torch.tensor(__lowerCamelCase       ,						dtype=torch.long					)					)

        if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
          lowerCamelCase_         =						flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

        self.assertListEqual(pt_generation_outputs.numpy().tolist()       ,						flax_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							List[str]					)     ->		Optional[int]:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						False
      lowerCamelCase_         =						max_length

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							Optional[int]					)     ->		Any:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						True
      lowerCamelCase_         =						max_length

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							int					)     ->		Dict:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						False
      lowerCamelCase_         =						max_length
      lowerCamelCase_         =						2

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							Any					)     ->		Union[str, Any]:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						False
      lowerCamelCase_         =						max_length
      lowerCamelCase_         =						2
      lowerCamelCase_         =						2

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[0]       ,						input_ids.shape[0] * config.num_return_sequences					)






    def    UpperCamelCase       (       self     :							Optional[Any]					)     ->		Dict:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						True
      lowerCamelCase_         =						max_length
      lowerCamelCase_         =						0.8
      lowerCamelCase_         =						10
      lowerCamelCase_         =						0.3
      lowerCamelCase_         =						1
      lowerCamelCase_         =						8
      lowerCamelCase_         =						9

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							Optional[Any]					)     ->		Optional[Any]:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						max_length
      lowerCamelCase_         =						1
      lowerCamelCase_         =						8
      lowerCamelCase_         =						9

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							List[str]					)     ->		Union[str, Any]:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()
      lowerCamelCase_         =						max_length
      lowerCamelCase_         =						2
      lowerCamelCase_         =						1
      lowerCamelCase_         =						8
      lowerCamelCase_         =						9

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							Union[str, Any]					)     ->		Dict:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()

      # pad attention mask on the left
      lowerCamelCase_         =						attention_mask.at[(0, 0)].set(0					)

      lowerCamelCase_         =						False
      lowerCamelCase_         =						max_length

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase       ,						attention_mask=__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase       ,						attention_mask=__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							Tuple					)     ->		Tuple:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()

      # pad attention mask on the left
      lowerCamelCase_         =						attention_mask.at[(0, 0)].set(0					)

      lowerCamelCase_         =						True
      lowerCamelCase_         =						max_length

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase       ,						attention_mask=__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase       ,						attention_mask=__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






    def    UpperCamelCase       (       self     :							Optional[int]					)     ->		List[Any]:
      lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						self._get_input_ids_and_config()

      # pad attention mask on the left
      lowerCamelCase_         =						attention_mask.at[(0, 0)].set(0					)

      lowerCamelCase_         =						2
      lowerCamelCase_         =						max_length

      for model_class in self.all_generative_model_classes:
        lowerCamelCase_         =						model_class(__lowerCamelCase					)

        lowerCamelCase_         =						model.generate(__lowerCamelCase       ,						attention_mask=__lowerCamelCase					).sequences
        self.assertEqual(generation_outputs.shape[-1]       ,						__lowerCamelCase					)

        lowerCamelCase_         =						jit(model.generate					)
        lowerCamelCase_         =						jit_generate(__lowerCamelCase       ,						attention_mask=__lowerCamelCase					).sequences

        self.assertListEqual(generation_outputs.tolist()       ,						jit_generation_outputs.tolist()					)






@require_flax
class      a						(					unittest.TestCase		):






    def    UpperCamelCase       (       self     :							Dict					)     ->		Optional[Any]:
      lowerCamelCase_         =						AutoTokenizer.from_pretrained('hf-internal-testing/tiny-bert'					)
      lowerCamelCase_         =						FlaxAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-bert-flax-only'					)

      lowerCamelCase_         =						'Hello world'
      lowerCamelCase_         =						tokenizer(__lowerCamelCase       ,						return_tensors='np'					).input_ids

      # typos are quickly detected (the correct argument is `do_sample`)
      with self.assertRaisesRegex(__lowerCamelCase       ,						'do_samples'					):
        model.generate(__lowerCamelCase       ,						do_samples=__lowerCamelCase					)

      # arbitrary arguments that will not be used anywhere are also not accepted
      with self.assertRaisesRegex(__lowerCamelCase       ,						'foo'					):
        lowerCamelCase_         =						{'foo': 'bar'}
        model.generate(__lowerCamelCase       ,						**__lowerCamelCase					)










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

import sentencepiece as spm

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


__lowerCAmelCase		:				str       			=		logging.get_logger(__name__)

__lowerCAmelCase		:				Dict       			=		{'vocab_file': 'spiece.model'}

__lowerCAmelCase		:				Optional[int]       			=		{
    'vocab_file': {
        'bert_for_seq_generation': (
            'https://huggingface.co/google/bert_for_seq_generation_L-24_bbc_encoder/resolve/main/spiece.model'
        ),
    }
}

__lowerCAmelCase		:				Dict       			=		{'bert_for_seq_generation': 512}




class     snake_case__							(_UpperCamelCase   ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_       :							Any							 =					VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE_       :							Union[str, Any]							 =					PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE_       :							int							 =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE_       :							List[int]							 =					[]
				SCREAMING_SNAKE_CASE_       :							Optional[Any]							 =					["""input_ids""", """attention_mask"""]





				def __init__(      self				:     List[Any]					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     str="<s>"					,      __lowerCamelCase				:     Optional[int]="</s>"					,      __lowerCamelCase				:     int="<unk>"					,      __lowerCamelCase				:     List[str]="<pad>"					,      __lowerCamelCase				:     Union[str, Any]="<::::>"					,      __lowerCamelCase				:     Optional[Dict[str, Any]] = None					,      **__lowerCamelCase				:     List[str]					,      )					->     None:
							a    =			{} if sp_model_kwargs is None else sp_model_kwargs

							# Add extra_ids to the special token list
							super().__init__(
							    bos_token=__lowerCamelCase					,      eos_token=__lowerCamelCase					,      unk_token=__lowerCamelCase					,      pad_token=__lowerCamelCase					,      sep_token=__lowerCamelCase					,      sp_model_kwargs=self.sp_model_kwargs					,      **__lowerCamelCase					,      )

							a    =			vocab_file

							a    =			spm.SentencePieceProcessor(**self.sp_model_kwargs							)
							self.sp_model.Load(__lowerCamelCase							)





				@property
				def 							__UpperCAmelCase							(      self				:     Dict							)					->     Dict:
							return self.sp_model.get_piece_size()





				def 							__UpperCAmelCase							(      self				:     Tuple							)					->     Optional[int]:
							a    =			{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				:     Dict							)					->     Optional[Any]:
							a    =			self.__dict__.copy()
							a    =			None
							return state





				def __setstate__(      self				:     Optional[Any]					,      __lowerCamelCase				:     Dict							)					->     Optional[Any]:
							a    =			d

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

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





				def 							__UpperCAmelCase							(      self				:     Union[str, Any]					,      __lowerCamelCase				:     str							)					->     List[str]:
							return self.sp_model.encode(__lowerCamelCase					,      out_type=__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     str					,      __lowerCamelCase				:     Union[str, Any]							)					->     int:
							return self.sp_model.piece_to_id(__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Tuple					,      __lowerCamelCase				:     List[str]							)					->     Any:
							a    =			self.sp_model.IdToPiece(__lowerCamelCase							)
							return token





				def 							__UpperCAmelCase							(      self				:     Any					,      __lowerCamelCase				:     Dict							)					->     Any:
							a    =			[]
							a    =			""
							for token in tokens:
										# make sure that special tokens are not decoded using sentencepiece model
										if token in self.all_special_tokens:
													out_string += self.sp_model.decode(__lowerCamelCase							) + token
													a    =			[]
										else:
													current_sub_tokens.append(__lowerCamelCase							)
							out_string += self.sp_model.decode(__lowerCamelCase							)
							return out_string.strip()





				def 							__UpperCAmelCase							(      self				:     Dict					,      __lowerCamelCase				:     str					,      __lowerCamelCase				:     Optional[str] = None							)					->     Tuple[str]:
							if not os.path.isdir(__lowerCamelCase							):
										logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""							)
										return
							a    =			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:
													a    =			self.sp_model.serialized_model_proto()
													fi.write(__lowerCamelCase							)

							return (out_vocab_file,)






