


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				:  List[Any]      					=							[
    {'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      snake_case_						( __lowercase=True       ):
   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=snake_case__						)						)
class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     Union[str, Any]			= None
    A_  :     Union[str, Any]			= None





    def 					_lowerCamelCase (	self				:   Tuple       ,   __snake_case				:   List[str]       ,   __snake_case				:   List[Any]	):





       '''simple docstring'''
       with TemporaryDirectory() as tmp_dir:
          UpperCAmelCase_							:      Optional[int]						=   dataset_module_factory(__snake_case       ,   cache_dir=__snake_case	)

          UpperCAmelCase_							:      List[Any]						=   import_main_class(dataset_module.module_path       ,   dataset=__snake_case	)

          UpperCAmelCase_							:      DatasetBuilder						=   builder_cls(
              cache_dir=__snake_case       ,   config_name=__snake_case       ,   hash=dataset_module.hash       ,   )

          UpperCAmelCase_							:      List[Any]						=   '''/'''.join(
              [
                  HF_GCP_BASE_URL,
                  builder_instance._relative_data_dir(with_hash=__snake_case	).replace(os.sep       ,   '''/'''	),
                  config.DATASET_INFO_FILENAME,
              ]	)
          UpperCAmelCase_							:      List[str]						=   cached_path(__snake_case       ,   cache_dir=__snake_case	)
          self.assertTrue(os.path.exists(__snake_case	)	)





@pytest.mark.integration
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Union[str, Any]						=   tmp_path_factory.mktemp('''test_hf_gcp'''       ) / '''test_wikipedia_simple'''
   UpperCAmelCase_							:      List[str]						=   dataset_module_factory('''wikipedia'''			,       cache_dir=__lowercase       )
   UpperCAmelCase_							:      int						=   import_main_class(dataset_module.module_path       )
   UpperCAmelCase_							:      DatasetBuilder						=   builder_cls(
       cache_dir=__lowercase			,       config_name='''20220301.frr'''			,       hash=dataset_module.hash			,       )
   # use the HF cloud storage, not the original download_and_prepare that uses apache-beam
   UpperCAmelCase_							:      List[Any]						=   None
   builder_instance.download_and_prepare()
   UpperCAmelCase_							:      Optional[int]						=   builder_instance.as_dataset()
   assert ds


@pytest.mark.integration
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      List[Any]						=   dataset_module_factory('''wikipedia'''			,       cache_dir=__lowercase       )
   UpperCAmelCase_							:      Dict						=   import_main_class(dataset_module.module_path			,       dataset=__lowercase       )
   UpperCAmelCase_							:      DatasetBuilder						=   builder_cls(
       cache_dir=__lowercase			,       config_name='''20220301.frr'''			,       hash=dataset_module.hash			,       )
   UpperCAmelCase_							:      List[Any]						=   builder_instance.as_streaming_dataset()
   assert ds
   assert isinstance(__lowercase			,       __lowercase       )
   assert "train" in ds
   assert isinstance(ds['''train''']			,       __lowercase       )
   assert next(iter(ds['''train''']       )       )



from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase				:  str      					=							logging.get_logger(__name__)

__UpperCamelCase				:  Any      					=							{
    'tiiuae/falcon-40b': 'https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json',
    'tiiuae/falcon-7b': 'https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json',
}





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     int			= 'falcon'
    A_  :     int			= ['past_key_values']





    def __init__(	self				:   Optional[Any]       ,   __snake_case				:   Tuple=65_024       ,   __snake_case				:   List[str]=4_544       ,   __snake_case				:   Optional[Any]=32       ,   __snake_case				:   Any=71       ,   __snake_case				:   str=1E-5       ,   __snake_case				:   List[str]=0.02       ,   __snake_case				:   List[Any]=True       ,   __snake_case				:   Dict=0.0       ,   __snake_case				:   Optional[Any]=0.0       ,   __snake_case				:   Any=None       ,   __snake_case				:   List[Any]=False       ,   __snake_case				:   Dict=False       ,   __snake_case				:   Optional[int]=True       ,   __snake_case				:   List[Any]=True       ,   __snake_case				:   Optional[Any]=False       ,   __snake_case				:   Dict=11       ,   __snake_case				:   List[str]=11       ,   **__snake_case				:   int       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   vocab_size
       # Backward compatibility with n_embed kwarg
       UpperCAmelCase_							:      Union[str, Any]						=   kwargs.pop('''n_embed'''       ,   __snake_case	)
       UpperCAmelCase_							:      str						=   hidden_size if n_embed is None else n_embed
       UpperCAmelCase_							:      Tuple						=   num_hidden_layers
       UpperCAmelCase_							:      Tuple						=   num_attention_heads
       UpperCAmelCase_							:      Optional[int]						=   layer_norm_epsilon
       UpperCAmelCase_							:      int						=   initializer_range
       UpperCAmelCase_							:      Optional[int]						=   use_cache
       UpperCAmelCase_							:      List[Any]						=   hidden_dropout
       UpperCAmelCase_							:      Any						=   attention_dropout

       UpperCAmelCase_							:      Tuple						=   bos_token_id
       UpperCAmelCase_							:      List[Any]						=   eos_token_id
       UpperCAmelCase_							:      Any						=   num_attention_heads if num_kv_heads is None else num_kv_heads
       UpperCAmelCase_							:      Optional[int]						=   alibi
       UpperCAmelCase_							:      Dict						=   new_decoder_architecture
       UpperCAmelCase_							:      List[Any]						=   multi_query  # Ignored when new_decoder_architecture is True
       UpperCAmelCase_							:      Tuple						=   parallel_attn
       UpperCAmelCase_							:      List[Any]						=   bias

       super().__init__(bos_token_id=__snake_case       ,   eos_token_id=__snake_case       ,   **__snake_case	)





    @property
    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       return self.hidden_size // self.num_attention_heads





    @property
    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       return not self.alibi