import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
					import jax.numpy as jnp

					from transformers.models.distilbert.modeling_flax_distilbert import (
					    FlaxDistilBertForMaskedLM,
					    FlaxDistilBertForMultipleChoice,
					    FlaxDistilBertForQuestionAnswering,
					    FlaxDistilBertForSequenceClassification,
					    FlaxDistilBertForTokenClassification,
					    FlaxDistilBertModel,
					)





class       lowercase__		(     unittest.TestCase):


						def __init__( self			:       Dict					,	UpperCamelCase__			:       Any					,	UpperCamelCase__			:       Optional[Any]=13					,	UpperCamelCase__			:       Optional[int]=7					,	UpperCamelCase__			:       str=True					,	UpperCamelCase__			:       Tuple=True					,	UpperCamelCase__			:       Tuple=True					,	UpperCamelCase__			:       Dict=True					,	UpperCamelCase__			:       Tuple=99					,	UpperCamelCase__			:       Tuple=32					,	UpperCamelCase__			:       List[str]=5					,	UpperCamelCase__			:       Tuple=4					,	UpperCamelCase__			:       Tuple=37					,	UpperCamelCase__			:       str="gelu"					,	UpperCamelCase__			:       List[Any]=0.1					,	UpperCamelCase__			:       Optional[Any]=0.1					,	UpperCamelCase__			:       Tuple=512					,	UpperCamelCase__			:       Optional[Any]=16					,	UpperCamelCase__			:       Any=2					,	UpperCamelCase__			:       int=0.02					,	UpperCamelCase__			:       Tuple=4					,	):





												'''simple docstring'''





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


						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							ids_tensor([self.batch_size, self.seq_length]					,	self.vocab_size		)

												SCREAMING_SNAKE_CASE							:					int           =							None
												if self.use_attention_mask:
																		SCREAMING_SNAKE_CASE							:					List[str]           =							random_attention_mask([self.batch_size, self.seq_length]		)

												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							DistilBertConfig(
												    vocab_size=self.vocab_size					,	dim=self.hidden_size					,	n_layers=self.num_hidden_layers					,	n_heads=self.num_attention_heads					,	hidden_dim=self.intermediate_size					,	hidden_act=self.hidden_act					,	dropout=self.hidden_dropout_prob					,	attention_dropout=self.attention_probs_dropout_prob					,	max_position_embeddings=self.max_position_embeddings					,	initializer_range=self.initializer_range					,	tie_weights_=UpperCamelCase__					,	)

												return config, input_ids, attention_mask



						def  __A ( self			:       int		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Any           =							self.prepare_config_and_inputs()
												SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE							:					Tuple           =							config_and_inputs
												SCREAMING_SNAKE_CASE							:					str           =							{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
												return config, inputs_dict






@require_flax
class       lowercase__		(     UpperCamelCase_			,							unittest.TestCase):
						UpperCamelCase_  =	(
						    (
						        FlaxDistilBertModel,
						        FlaxDistilBertForMaskedLM,
						        FlaxDistilBertForMultipleChoice,
						        FlaxDistilBertForQuestionAnswering,
						        FlaxDistilBertForSequenceClassification,
						        FlaxDistilBertForTokenClassification,
						        FlaxDistilBertForQuestionAnswering,
						    )
						    if is_flax_available()
						    else ()
						)


						def  __A ( self			:       Any		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							FlaxDistilBertModelTester(self		)



						@slow
						def  __A ( self			:       Any		):





												'''simple docstring'''





												for model_class_name in self.all_model_classes:
																		SCREAMING_SNAKE_CASE							:					str           =							model_class_name.from_pretrained('''distilbert-base-uncased'''		)
																		SCREAMING_SNAKE_CASE							:					Dict           =							model(np.ones((1, 1)		)		)
																		self.assertIsNotNone(UpperCamelCase__		)






@require_flax
class       lowercase__		(     unittest.TestCase):


						@slow
						def  __A ( self			:       int		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Any           =							FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased'''		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]		)
												SCREAMING_SNAKE_CASE							:					Any           =							np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
												SCREAMING_SNAKE_CASE							:					Tuple           =							model(UpperCamelCase__					,	attention_mask=UpperCamelCase__		)[0]
												SCREAMING_SNAKE_CASE							:					Tuple           =							(1, 11, 768)
												self.assertEqual(output.shape					,	UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					str           =							np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]		)

												self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]					,	UpperCamelCase__					,	atol=1E-4		)		)






import argparse
import json
import logging
import os
import sys
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


__UpperCamelCase	: str    		=  [
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        'text-classification',
        'language-modeling',
        'summarization',
        'token-classification',
        'question-answering',
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
					import run_clm_flax
					import run_flax_glue
					import run_flax_ner
					import run_mlm_flax
					import run_qa
					import run_summarization_flax
					import run_ta_mlm_flax


logging.basicConfig(level=logging.DEBUG)

__UpperCamelCase	: int    		=  logging.getLogger()
def    A			(					):
						SCREAMING_SNAKE_CASE							:					str           =							argparse.ArgumentParser()
						parser.add_argument('''-f'''		)
						SCREAMING_SNAKE_CASE							:					List[str]           =							parser.parse_args()
						return args.f
def    A			(					_lowercase	,						_lowercase="eval"		):
						SCREAMING_SNAKE_CASE							:					Dict           =							os.path.join(_lowercase	,						f"""{split}_results.json"""		)
						if os.path.exists(_lowercase		):
												with open(_lowercase	,						'''r'''		) as f:
																		return json.load(_lowercase		)
						raise ValueError(f"""can't find {path}"""		)


__UpperCamelCase	: Optional[Any]    		=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)





class       lowercase__		(     UpperCamelCase_):


						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					int           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Tuple           =							f"""
            run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --eval_steps=2
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_glue.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					str           =							f"""
            run_clm_flax.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --block_size 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_clm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Dict           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	100		)


						@slow
						def  __A ( self			:       Any		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_summarization.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --test_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=8
            --do_train
            --do_eval
            --do_predict
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --predict_with_generate
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_summarization_flax.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__					,	split='''test'''		)
																		self.assertGreaterEqual(result['''test_rouge1''']					,	10		)
																		self.assertGreaterEqual(result['''test_rouge2''']					,	2		)
																		self.assertGreaterEqual(result['''test_rougeL''']					,	7		)
																		self.assertGreaterEqual(result['''test_rougeLsum''']					,	7		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Dict           =							f"""
            run_mlm.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --logging_steps 2 --eval_steps 2
            --do_train
            --do_eval
            --num_train_epochs=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	42		)


						@slow
						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_t5_mlm_flax.py
            --model_name_or_path t5-small
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_ta_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.42		)


						@slow
						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Tuple           =							7 if get_gpu_count() > 1 else 2

												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Any           =							f"""
            run_flax_ner.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --warmup_steps=2
            --learning_rate=2e-4
            --logging_steps 2 --eval_steps 2
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_ner.main()
																		SCREAMING_SNAKE_CASE							:					List[str]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	0.3		)



						@slow
						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_qa.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=2
            --do_train
            --do_eval
            --logging_steps 2 --eval_steps 2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_qa.main()
																		SCREAMING_SNAKE_CASE							:					str           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	30		)
																		self.assertGreaterEqual(result['''eval_exact''']					,	30		)





