
'''simple docstring'''






import os
from tempfile import TemporaryDirectory
from unittest import TestCase

import pytest
from absl.testing import parameterized

from datasets import config
from datasets.arrow_reader import HF_GCP_BASE_URL
from datasets.builder import DatasetBuilder
from datasets.dataset_dict import IterableDatasetDict
from datasets.iterable_dataset import IterableDataset
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.file_utils import cached_path


UpperCAmelCase_							: Tuple =							[
    {'dataset': 'wikipedia', 'config_name': '20220301.de'},
    {'dataset': 'wikipedia', 'config_name': '20220301.en'},
    {'dataset': 'wikipedia', 'config_name': '20220301.fr'},
    {'dataset': 'wikipedia', 'config_name': '20220301.frr'},
    {'dataset': 'wikipedia', 'config_name': '20220301.it'},
    {'dataset': 'wikipedia', 'config_name': '20220301.simple'},
    {'dataset': 'snli', 'config_name': 'plain_text'},
    {'dataset': 'eli5', 'config_name': 'LFQA_reddit'},
    {'dataset': 'wiki40b', 'config_name': 'en'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.nq.compressed'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.nq.no_index'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.multiset.no_index'},
    {'dataset': 'natural_questions', 'config_name': 'default'},
]







def 		A_							(						_lowerCAmelCase :     Optional[Any]=True				):



      """simple docstring"""




      if with_config:
            return [
                {
                    "testcase_name": d["dataset"] + "/" + d["config_name"],
                    "dataset": d["dataset"],
                    "config_name": d["config_name"],
                }
                for d in DATASETS_ON_HF_GCP
            ]
      else:
            return [
                {"testcase_name": dataset, "dataset": dataset} for dataset in {d["dataset"] for d in DATASETS_ON_HF_GCP}
            ]

@parameterized.named_parameters(list_datasets_on_hf_gcp_parameters(with_config=A			)			)
class 			UpperCAmelCase__    (							A			):
    lowerCAmelCase_					     =		None
    lowerCAmelCase_					     =		None



    def 	lowerCamelCase_	(						self   :							Union[str, Any],__A   :							str,__A   :							int  ):
          with TemporaryDirectory() as tmp_dir:
                _lowerCamelCase					: Union[str, Any]		    = dataset_module_factory(__A,cache_dir=__A  )

                _lowerCamelCase					: Any		    = import_main_class(dataset_module.module_path,dataset=__A  )

                _lowerCamelCase					: DatasetBuilder		    = builder_cls(
                    cache_dir=__A,config_name=__A,hash=dataset_module.hash,)

                _lowerCamelCase					: List[Any]		    = "/".join(
                    [
                        HF_GCP_BASE_URL,
                        builder_instance._relative_data_dir(with_hash=__A  ).replace(os.sep,"/"  ),
                        config.DATASET_INFO_FILENAME,
                    ]  )
                _lowerCamelCase					: Any		    = cached_path(__A,cache_dir=__A  )
                self.assertTrue(os.path.exists(__A  )  )







@pytest.mark.integration
def 		A_							(						_lowerCAmelCase :     Tuple				):



      """simple docstring"""




      _lowerCamelCase					: Tuple		    = tmp_path_factory.mktemp("test_hf_gcp"				) / "test_wikipedia_simple"
      _lowerCamelCase					: str		    = dataset_module_factory("wikipedia"					,						cache_dir=_lowerCAmelCase				)
      _lowerCamelCase					: List[str]		    = import_main_class(dataset_module.module_path				)
      _lowerCamelCase					: DatasetBuilder		    = builder_cls(
          cache_dir=_lowerCAmelCase					,						config_name="20220301.frr"					,						hash=dataset_module.hash					,						)
      # use the HF cloud storage, not the original download_and_prepare that uses apache-beam
      _lowerCamelCase					: List[str]		    = None
      builder_instance.download_and_prepare()
      _lowerCamelCase					: str		    = builder_instance.as_dataset()
      assert ds







@pytest.mark.integration
def 		A_							(						_lowerCAmelCase :     str				):



      """simple docstring"""




      _lowerCamelCase					: List[Any]		    = dataset_module_factory("wikipedia"					,						cache_dir=_lowerCAmelCase				)
      _lowerCamelCase					: Optional[int]		    = import_main_class(dataset_module.module_path					,						dataset=_lowerCAmelCase				)
      _lowerCamelCase					: DatasetBuilder		    = builder_cls(
          cache_dir=_lowerCAmelCase					,						config_name="20220301.frr"					,						hash=dataset_module.hash					,						)
      _lowerCamelCase					: str		    = builder_instance.as_streaming_dataset()
      assert ds
      assert isinstance(_lowerCAmelCase					,						_lowerCAmelCase				)
      assert "train" in ds
      assert isinstance(ds["train"]					,						_lowerCAmelCase				)
      assert next(iter(ds["train"]				)				)
import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
  from transformers.models.bert.modeling_flax_bert import (
      FlaxBertForMaskedLM,
      FlaxBertForMultipleChoice,
      FlaxBertForNextSentencePrediction,
      FlaxBertForPreTraining,
      FlaxBertForQuestionAnswering,
      FlaxBertForSequenceClassification,
      FlaxBertForTokenClassification,
      FlaxBertModel,
  )




class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




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



    """simple docstring"""



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



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



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

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

    _UpperCamelCase:  Optional[int]    						=			None
    if self.use_token_type_ids:
     _UpperCamelCase:  Tuple    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.type_vocab_size				)

    _UpperCamelCase:  str    						=			BertConfig(
        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=_lowercase		,	initializer_range=self.initializer_range		,	)

    return config, input_ids, token_type_ids, attention_mask



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			self.prepare_config_and_inputs()
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  Any    						=			config_and_inputs
    _UpperCamelCase:  List[Any]    						=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
    return config, inputs_dict



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			self.prepare_config_and_inputs()
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  str    						=			config_and_inputs

    _UpperCamelCase:  int    						=			True
    _UpperCamelCase:  Optional[Any]    						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
    _UpperCamelCase:  Optional[Any]    						=			ids_tensor([self.batch_size, self.seq_length]		,	vocab_size=2				)

    return (
        config,
        input_ids,
        attention_mask,
        encoder_hidden_states,
        encoder_attention_mask,
    )




@require_flax
class 			__magic_name__	(		__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Dict	      = True

   lowerCAmelCase   :     List[Any]	      = (
       (
           FlaxBertModel,
           FlaxBertForPreTraining,
           FlaxBertForMaskedLM,
           FlaxBertForMultipleChoice,
           FlaxBertForQuestionAnswering,
           FlaxBertForNextSentencePrediction,
           FlaxBertForSequenceClassification,
           FlaxBertForTokenClassification,
           FlaxBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			FlaxBertModelTester(self				)



   @slow
   def 	lowerCAmelCase    (						self       :	Tuple				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			FlaxBertModel.from_pretrained('''bert-base-cased'''				)
    _UpperCamelCase:  int    						=			model(np.ones((1, 1)				)				)
    self.assertIsNotNone(_lowercase				)