"""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

				_UpperCamelCase								=		"""0.12"""  # assumed parallelism: 8


if is_torch_available():
				import torch

def  _A(		lowerCAmelCase			,  lowerCAmelCase			,  lowerCAmelCase=None							):







							if rng is None:
														A__       :     str								=			random.Random()

							A__       :     List[Any]								=			1
							for dim in shape:
														total_dims *= dim

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

							A__       :     int								=			np.array(_A			,  dtype=jnp.intaa							).reshape(_A							)

							return output

def  _A(		lowerCAmelCase			,  lowerCAmelCase=None							):
							A__       :     Optional[int]								=			ids_tensor(_A			,  vocab_size=2			,  rng=_A							)
							# make sure that at least one token is attended to for each batch
							A__       :     Optional[Any]								=			1
							return attn_mask
@require_flax
class 						__UpperCAmelCase    :







			'''simple docstring'''
			_UpperCamelCase  :						str					     =				None
			_UpperCamelCase  :						List[Any]					     =				()





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[str]								=			self.model_tester.prepare_config_and_inputs_for_common()

										# cut to half length & take max batch_size 3
										A__       :     int								=			2
										A__       :     List[str]								=			inputs['input_ids'].shape[-1] // 2
										A__       :     Optional[Any]								=			inputs['input_ids'][:max_batch_size, :sequence_length]

										A__       :     Optional[int]								=			jnp.ones_like(snake_case__							)
										A__       :     List[Any]								=			attention_mask[:max_batch_size, :sequence_length]

										# generate max 5 tokens
										A__       :     Tuple								=			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()`
																	A__       :     Any								=			config.eos_token_id
										return config, input_ids, attention_mask, max_length





			@is_pt_flax_cross_test
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Any								=			self._get_input_ids_and_config()
										A__       :     List[Any]								=			False
										A__       :     Optional[int]								=			max_length
										A__       :     List[str]								=			0

										for model_class in self.all_generative_model_classes:
																	A__       :     Dict								=			model_class(snake_case__							)

																	A__       :     Optional[Any]								=			model_class.__name__[4:]  # Skip the "Flax" at the beginning
																	A__       :     Tuple								=			getattr(snake_case__   ,      snake_case__							)
																	A__       :     Optional[int]								=			pt_model_class(snake_case__							).eval()
																	A__       :     int								=			load_flax_weights_in_pytorch_model(snake_case__   ,      flax_model.params							)

																	A__       :     int								=			flax_model.generate(snake_case__							).sequences
																	A__       :     str								=			pt_model.generate(torch.tensor(snake_case__   ,      dtype=torch.long							)							)

																	if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																								A__       :     List[Any]								=			flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Dict								=			self._get_input_ids_and_config()
										A__       :     Optional[Any]								=			False
										A__       :     str								=			max_length

										for model_class in self.all_generative_model_classes:
																	A__       :     Optional[Any]								=			model_class(snake_case__							)

																	A__       :     List[Any]								=			model.generate(snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     int								=			jit(model.generate							)
																	A__       :     Optional[int]								=			jit_generate(snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     str								=			self._get_input_ids_and_config()
										A__       :     Union[str, Any]								=			True
										A__       :     int								=			max_length

										for model_class in self.all_generative_model_classes:
																	A__       :     Optional[Any]								=			model_class(snake_case__							)

																	A__       :     List[str]								=			model.generate(snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     Tuple								=			jit(model.generate							)
																	A__       :     int								=			jit_generate(snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Dict								=			self._get_input_ids_and_config()
										A__       :     Dict								=			False
										A__       :     str								=			max_length
										A__       :     Dict								=			2

										for model_class in self.all_generative_model_classes:
																	A__       :     Optional[int]								=			model_class(snake_case__							)

																	A__       :     Optional[Any]								=			model.generate(snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     Optional[Any]								=			jit(model.generate							)
																	A__       :     List[str]								=			jit_generate(snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[str]								=			self._get_input_ids_and_config()
										A__       :     Any								=			False
										A__       :     Any								=			max_length
										A__       :     Any								=			2
										A__       :     List[Any]								=			2

										for model_class in self.all_generative_model_classes:
																	A__       :     Tuple								=			model_class(snake_case__							)

																	A__       :     str								=			model.generate(snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[0]   ,      input_ids.shape[0] * config.num_return_sequences							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     str								=			self._get_input_ids_and_config()
										A__       :     List[Any]								=			True
										A__       :     int								=			max_length
										A__       :     List[str]								=			0.8
										A__       :     Union[str, Any]								=			10
										A__       :     Any								=			0.3
										A__       :     List[Any]								=			1
										A__       :     int								=			8
										A__       :     List[Any]								=			9

										for model_class in self.all_generative_model_classes:
																	A__       :     Optional[int]								=			model_class(snake_case__							)

																	A__       :     str								=			model.generate(snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     str								=			jit(model.generate							)
																	A__       :     str								=			jit_generate(snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Union[str, Any]								=			self._get_input_ids_and_config()
										A__       :     Dict								=			max_length
										A__       :     str								=			1
										A__       :     Union[str, Any]								=			8
										A__       :     int								=			9

										for model_class in self.all_generative_model_classes:
																	A__       :     str								=			model_class(snake_case__							)

																	A__       :     List[Any]								=			model.generate(snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     List[str]								=			jit(model.generate							)
																	A__       :     Optional[Any]								=			jit_generate(snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Tuple								=			self._get_input_ids_and_config()
										A__       :     str								=			max_length
										A__       :     str								=			2
										A__       :     str								=			1
										A__       :     List[str]								=			8
										A__       :     Union[str, Any]								=			9

										for model_class in self.all_generative_model_classes:
																	A__       :     Tuple								=			model_class(snake_case__							)

																	A__       :     int								=			model.generate(snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     Dict								=			jit(model.generate							)
																	A__       :     Union[str, Any]								=			jit_generate(snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Any								=			self._get_input_ids_and_config()

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

										A__       :     Dict								=			False
										A__       :     Tuple								=			max_length

										for model_class in self.all_generative_model_classes:
																	A__       :     Optional[Any]								=			model_class(snake_case__							)

																	A__       :     Dict								=			model.generate(snake_case__   ,      attention_mask=snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     List[str]								=			jit(model.generate							)
																	A__       :     List[Any]								=			jit_generate(snake_case__   ,      attention_mask=snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[Any]								=			self._get_input_ids_and_config()

										# pad attention mask on the left
										A__       :     List[Any]								=			attention_mask.at[(0, 0)].set(0							)

										A__       :     Union[str, Any]								=			True
										A__       :     Dict								=			max_length

										for model_class in self.all_generative_model_classes:
																	A__       :     Optional[int]								=			model_class(snake_case__							)

																	A__       :     List[Any]								=			model.generate(snake_case__   ,      attention_mask=snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     Dict								=			jit(model.generate							)
																	A__       :     List[str]								=			jit_generate(snake_case__   ,      attention_mask=snake_case__							).sequences

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





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Optional[int]								=			self._get_input_ids_and_config()

										# pad attention mask on the left
										A__       :     List[str]								=			attention_mask.at[(0, 0)].set(0							)

										A__       :     Optional[Any]								=			2
										A__       :     Dict								=			max_length

										for model_class in self.all_generative_model_classes:
																	A__       :     Optional[int]								=			model_class(snake_case__							)

																	A__       :     Union[str, Any]								=			model.generate(snake_case__   ,      attention_mask=snake_case__							).sequences
																	self.assertEqual(generation_outputs.shape[-1]   ,      snake_case__							)

																	A__       :     str								=			jit(model.generate							)
																	A__       :     int								=			jit_generate(snake_case__   ,      attention_mask=snake_case__							).sequences

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





@require_flax
class 						__UpperCAmelCase    (unittest.TestCase     ):







			'''simple docstring'''





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[Any]								=			AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-bert"""							)
										A__       :     Optional[Any]								=			FlaxAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""							)

										A__       :     Any								=			'Hello world'
										A__       :     Any								=			tokenizer(snake_case__   ,      return_tensors="""np"""							).input_ids

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

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









'''simple docstring'''







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

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_gpta import GPTaTokenizer


if TYPE_CHECKING:
       from transformers.pipelines.conversational import Conversation


lowerCAmelCase     :  Tuple											=    logging.get_logger(__name__)

lowerCAmelCase     :  Union[str, Any]											=    {"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

lowerCAmelCase     :  str											=    {
    """vocab_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/vocab.json""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/vocab.json""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/vocab.json""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/vocab.json""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/merges.txt""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/merges.txt""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/merges.txt""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/merges.txt""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/tokenizer.json""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/tokenizer.json""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/tokenizer.json""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/tokenizer.json""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/tokenizer.json""",
    },
}

lowerCAmelCase     :  str											=    {
    """gpt2""": 10_24,
    """gpt2-medium""": 10_24,
    """gpt2-large""": 10_24,
    """gpt2-xl""": 10_24,
    """distilgpt2""": 10_24,
}
class 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_						):



     """simple docstring"""






     __magic_name__           =     VOCAB_FILES_NAMES
     __magic_name__           =     PRETRAINED_VOCAB_FILES_MAP
     __magic_name__           =     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __magic_name__           =     ["input_ids", "attention_mask"]
     __magic_name__           =     GPTaTokenizer







     def __init__(			self   ,						snake_case__=None   ,						snake_case__=None   ,						snake_case__=None   ,						snake_case__="<|endoftext|>"   ,						snake_case__="<|endoftext|>"   ,						snake_case__="<|endoftext|>"   ,						snake_case__=False   ,						**snake_case__   ,						):

            '''simple docstring'''

            super().__init__(
                snake_case__   ,						snake_case__   ,						tokenizer_file=snake_case__   ,						unk_token=snake_case__   ,						bos_token=snake_case__   ,						eos_token=snake_case__   ,						add_prefix_space=snake_case__   ,						**snake_case__   ,						)

            _lowerCAmelCase     :				str       =							kwargs.pop('add_bos_token'   ,						snake_case__ )

            _lowerCAmelCase     :				List[Any]       =							json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__() )
            if pre_tok_state.get('add_prefix_space'   ,						snake_case__ ) != add_prefix_space:
                   _lowerCAmelCase     :				str       =							getattr(snake_case__   ,						pre_tok_state.pop('type' ) )
                   _lowerCAmelCase     :				List[str]       =							add_prefix_space
                   _lowerCAmelCase     :				Dict       =							pre_tok_class(**snake_case__ )

            _lowerCAmelCase     :				Any       =							add_prefix_space







     def a						(			self   ,						*snake_case__   ,						**snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				int       =							kwargs.get('is_split_into_words'   ,						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(*snake_case__   ,						**snake_case__ )







     def a						(			self   ,						*snake_case__   ,						**snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				int       =							kwargs.get('is_split_into_words'   ,						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(*snake_case__   ,						**snake_case__ )







     def a						(			self   ,						snake_case__   ,						snake_case__ = None ):

            '''simple docstring'''

            _lowerCAmelCase     :				Optional[int]       =							self._tokenizer.model.save(snake_case__   ,						name=snake_case__ )
            return tuple(snake_case__ )







     def a						(			self   ,						snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				Any       =							[]
            for is_user, text in conversation.iter_texts():
                   input_ids.extend(self.encode(snake_case__   ,						add_special_tokens=snake_case__ ) + [self.eos_token_id] )

            if len(snake_case__ ) > self.model_max_length:
                   _lowerCAmelCase     :				str       =							input_ids[-self.model_max_length :]
            return input_ids

