



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

							__lowerCAmelCase						      =       '''0.12'''  # assumed parallelism: 8


if is_torch_available():
							import torch



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_=None  )      ->				Union[str, Any]:
		if rng is None:
				_a					:				List[str]				=							random.Random()

		_a					:				Optional[Any]				=							1
		for dim in shape:
				total_dims *= dim

		_a					:				Optional[int]				=							[]
		for _ in range(lowerCAmelCase_  ):
				values.append(rng.randint(0						,      vocab_size - 1  )  )

		_a					:				Any				=							np.array(lowerCAmelCase_						,      dtype=jnp.intaa  ).reshape(lowerCAmelCase_  )

		return output



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_=None  )      ->				Optional[Any]:
		_a					:				Union[str, Any]				=							ids_tensor(lowerCAmelCase_						,      vocab_size=2						,      rng=lowerCAmelCase_  )
		# make sure that at least one token is attended to for each batch
		_a					:				int				=							1
		return attn_mask

@require_flax
class      __magic_name__						:
							lowerCAmelCase	:   Any           = None
							lowerCAmelCase	:   List[Any]           = ()







							def 							__lowercase      (							self :	Union[str, Any]		):
									_a					:				List[str]				=							self.model_tester.prepare_config_and_inputs_for_common()

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

									_a					:				Dict				=							jnp.ones_like(_UpperCAmelCase		)
									_a					:				List[Any]				=							attention_mask[:max_batch_size, :sequence_length]

									# generate max 5 tokens
									_a					:				Union[str, Any]				=							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					:				List[str]				=							config.eos_token_id
									return config, input_ids, attention_mask, max_length







							@is_pt_flax_cross_test
							def 							__lowercase      (							self :	Optional[int]		):
									_a					:				Union[str, Any]				=							self._get_input_ids_and_config()
									_a					:				Union[str, Any]				=							False
									_a					:				List[Any]				=							max_length
									_a					:				Tuple				=							0

									for model_class in self.all_generative_model_classes:
											_a					:				Dict				=							model_class(_UpperCAmelCase		)

											_a					:				List[str]				=							model_class.__name__[4:]  # Skip the "Flax" at the beginning
											_a					:				Tuple				=							getattr(_UpperCAmelCase     ,_UpperCAmelCase		)
											_a					:				str				=							pt_model_class(_UpperCAmelCase		).eval()
											_a					:				int				=							load_flax_weights_in_pytorch_model(_UpperCAmelCase     ,flax_model.params		)

											_a					:				List[Any]				=							flax_model.generate(_UpperCAmelCase		).sequences
											_a					:				Optional[Any]				=							pt_model.generate(torch.tensor(_UpperCAmelCase     ,dtype=torch.long		)		)

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

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







							def 							__lowercase      (							self :	Any		):
									_a					:				Optional[Any]				=							self._get_input_ids_and_config()
									_a					:				Tuple				=							False
									_a					:				List[str]				=							max_length

									for model_class in self.all_generative_model_classes:
											_a					:				str				=							model_class(_UpperCAmelCase		)

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

											_a					:				int				=							jit(model.generate		)
											_a					:				List[Any]				=							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	Union[str, Any]		):
									_a					:				Any				=							self._get_input_ids_and_config()
									_a					:				Tuple				=							True
									_a					:				List[Any]				=							max_length

									for model_class in self.all_generative_model_classes:
											_a					:				List[str]				=							model_class(_UpperCAmelCase		)

											_a					:				Tuple				=							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]     ,_UpperCAmelCase		)

											_a					:				Optional[Any]				=							jit(model.generate		)
											_a					:				Any				=							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	int		):
									_a					:				List[Any]				=							self._get_input_ids_and_config()
									_a					:				Any				=							False
									_a					:				str				=							max_length
									_a					:				Dict				=							2

									for model_class in self.all_generative_model_classes:
											_a					:				List[Any]				=							model_class(_UpperCAmelCase		)

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

											_a					:				int				=							jit(model.generate		)
											_a					:				Tuple				=							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	Dict		):
									_a					:				str				=							self._get_input_ids_and_config()
									_a					:				int				=							False
									_a					:				Optional[Any]				=							max_length
									_a					:				Union[str, Any]				=							2
									_a					:				str				=							2

									for model_class in self.all_generative_model_classes:
											_a					:				Any				=							model_class(_UpperCAmelCase		)

											_a					:				int				=							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[0]     ,input_ids.shape[0] * config.num_return_sequences		)







							def 							__lowercase      (							self :	Optional[int]		):
									_a					:				Tuple				=							self._get_input_ids_and_config()
									_a					:				Optional[Any]				=							True
									_a					:				int				=							max_length
									_a					:				List[Any]				=							0.8
									_a					:				Dict				=							10
									_a					:				Union[str, Any]				=							0.3
									_a					:				Tuple				=							1
									_a					:				List[str]				=							8
									_a					:				Union[str, Any]				=							9

									for model_class in self.all_generative_model_classes:
											_a					:				Dict				=							model_class(_UpperCAmelCase		)

											_a					:				Tuple				=							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]     ,_UpperCAmelCase		)

											_a					:				List[str]				=							jit(model.generate		)
											_a					:				Tuple				=							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	Union[str, Any]		):
									_a					:				Tuple				=							self._get_input_ids_and_config()
									_a					:				int				=							max_length
									_a					:				Dict				=							1
									_a					:				Union[str, Any]				=							8
									_a					:				int				=							9

									for model_class in self.all_generative_model_classes:
											_a					:				str				=							model_class(_UpperCAmelCase		)

											_a					:				Tuple				=							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]     ,_UpperCAmelCase		)

											_a					:				Any				=							jit(model.generate		)
											_a					:				Optional[Any]				=							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	Tuple		):
									_a					:				Dict				=							self._get_input_ids_and_config()
									_a					:				Optional[int]				=							max_length
									_a					:				Tuple				=							2
									_a					:				Tuple				=							1
									_a					:				int				=							8
									_a					:				int				=							9

									for model_class in self.all_generative_model_classes:
											_a					:				Union[str, Any]				=							model_class(_UpperCAmelCase		)

											_a					:				Dict				=							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]     ,_UpperCAmelCase		)

											_a					:				Optional[Any]				=							jit(model.generate		)
											_a					:				Dict				=							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	Dict		):
									_a					:				Any				=							self._get_input_ids_and_config()

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

									_a					:				Optional[int]				=							False
									_a					:				int				=							max_length

									for model_class in self.all_generative_model_classes:
											_a					:				str				=							model_class(_UpperCAmelCase		)

											_a					:				List[str]				=							model.generate(_UpperCAmelCase     ,attention_mask=_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]     ,_UpperCAmelCase		)

											_a					:				Dict				=							jit(model.generate		)
											_a					:				Optional[int]				=							jit_generate(_UpperCAmelCase     ,attention_mask=_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	Tuple		):
									_a					:				str				=							self._get_input_ids_and_config()

									# pad attention mask on the left
									_a					:				Optional[int]				=							attention_mask.at[(0, 0)].set(0		)

									_a					:				Optional[int]				=							True
									_a					:				List[Any]				=							max_length

									for model_class in self.all_generative_model_classes:
											_a					:				Tuple				=							model_class(_UpperCAmelCase		)

											_a					:				int				=							model.generate(_UpperCAmelCase     ,attention_mask=_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]     ,_UpperCAmelCase		)

											_a					:				Any				=							jit(model.generate		)
											_a					:				Dict				=							jit_generate(_UpperCAmelCase     ,attention_mask=_UpperCAmelCase		).sequences

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







							def 							__lowercase      (							self :	Optional[int]		):
									_a					:				Dict				=							self._get_input_ids_and_config()

									# pad attention mask on the left
									_a					:				Optional[int]				=							attention_mask.at[(0, 0)].set(0		)

									_a					:				Dict				=							2
									_a					:				Tuple				=							max_length

									for model_class in self.all_generative_model_classes:
											_a					:				List[str]				=							model_class(_UpperCAmelCase		)

											_a					:				Optional[int]				=							model.generate(_UpperCAmelCase     ,attention_mask=_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]     ,_UpperCAmelCase		)

											_a					:				Optional[Any]				=							jit(model.generate		)
											_a					:				Any				=							jit_generate(_UpperCAmelCase     ,attention_mask=_UpperCAmelCase		).sequences

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

@require_flax
class      __magic_name__						(   unittest.TestCase       ):







							def 							__lowercase      (							self :	int		):
									_a					:				Union[str, Any]				=							AutoTokenizer.from_pretrained('hf-internal-testing/tiny-bert'		)
									_a					:				str				=							FlaxAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-bert-flax-only'		)

									_a					:				Any				=							'Hello world'
									_a					:				Union[str, Any]				=							tokenizer(_UpperCAmelCase     ,return_tensors='np'		).input_ids

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

									# arbitrary arguments that will not be used anywhere are also not accepted
									with self.assertRaisesRegex(_UpperCAmelCase     ,'foo'		):
											_a					:				List[str]				=							{'foo': 'bar'}
											model.generate(_UpperCAmelCase     ,**_UpperCAmelCase		)









'''simple docstring'''

from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__lowerCAmelCase						      =       logging.get_logger(__name__)

__lowerCAmelCase						      =       {
    '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/config.json''',
    '''distilbert-base-uncased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/config.json''',
    '''distilbert-base-cased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-german-cased''': '''https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json''',
    '''distilbert-base-multilingual-cased''': (
        '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json'''
    ),
    '''distilbert-base-uncased-finetuned-sst-2-english''': (
        '''https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json'''
    ),
}

class      __magic_name__						(   _UpperCamelCase       ):
							lowerCAmelCase	:   Optional[int]           = 'distilbert'
							lowerCAmelCase	:   Optional[int]           = {
							    'hidden_size': 'dim',
							    'num_attention_heads': 'n_heads',
							    'num_hidden_layers': 'n_layers',
							}







							def __init__(							self :	str     ,_UpperCAmelCase :	Optional[Any]=30522     ,_UpperCAmelCase :	int=512     ,_UpperCAmelCase :	int=False     ,_UpperCAmelCase :	str=6     ,_UpperCAmelCase :	int=12     ,_UpperCAmelCase :	Union[str, Any]=768     ,_UpperCAmelCase :	Optional[Any]=4 * 768     ,_UpperCAmelCase :	int=0.1     ,_UpperCAmelCase :	Union[str, Any]=0.1     ,_UpperCAmelCase :	Tuple="gelu"     ,_UpperCAmelCase :	Any=0.02     ,_UpperCAmelCase :	List[Any]=0.1     ,_UpperCAmelCase :	Dict=0.2     ,_UpperCAmelCase :	Union[str, Any]=0     ,**_UpperCAmelCase :	List[Any]     ,):
									_a					:				Optional[Any]				=							vocab_size
									_a					:				Dict				=							max_position_embeddings
									_a					:				Optional[int]				=							sinusoidal_pos_embds
									_a					:				List[Any]				=							n_layers
									_a					:				Any				=							n_heads
									_a					:				int				=							dim
									_a					:				Optional[Any]				=							hidden_dim
									_a					:				str				=							dropout
									_a					:				Optional[Any]				=							attention_dropout
									_a					:				Any				=							activation
									_a					:				Tuple				=							initializer_range
									_a					:				Union[str, Any]				=							qa_dropout
									_a					:				List[Any]				=							seq_classif_dropout
									super().__init__(**_UpperCAmelCase     ,pad_token_id=_UpperCAmelCase		)

class      __magic_name__						(   _UpperCamelCase       ):







							@property
							def 							__lowercase      (							self :	Union[str, Any]		):
									if self.task == "multiple-choice":
											_a					:				Union[str, Any]				=							{0: 'batch', 1: 'choice', 2: 'sequence'}
									else:
											_a					:				List[str]				=							{0: 'batch', 1: 'sequence'}
									return OrderedDict(
									    [
									        ('input_ids', dynamic_axis),
									        ('attention_mask', dynamic_axis),
									    ]		)




