





import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

UpperCAmelCase_ 							=       [
    """cross_validation.py""",
    """gradient_accumulation.py""",
    """local_sgd.py""",
    """multi_process_metrics.py""",
    """memory.py""",
    """automatic_gradient_accumulation.py""",
    """fsdp_with_peak_mem_tracking.py""",
    """deepspeed_with_config_support.py""",
    """megatron_lm_gpt_pretraining.py""",
]
class      UpperCamelCase_						(					unittest.TestCase   ):




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_ = None			,  lowerCAmelCase_ = None					)						->							List[str]:
										_snake_case													=			None
										_snake_case													=			os.path.abspath(os.path.join('examples'			,  'by_feature'					)					)
										_snake_case													=			os.path.abspath('examples'					)
										for item in os.listdir(lowerCAmelCase_					):
															if item not in EXCLUDE_EXAMPLES:
																				_snake_case													=			os.path.join(lowerCAmelCase_			,  lowerCAmelCase_					)
																				if os.path.isfile(lowerCAmelCase_					) and ".py" in item_path:
																									with self.subTest(
																									    tested_script=lowerCAmelCase_			,  feature_script=lowerCAmelCase_			,  tested_section='main()' if parser_only else 'training_function()'			,  ):
																														_snake_case													=			compare_against_test(
																														    os.path.join(lowerCAmelCase_			,  lowerCAmelCase_					)			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)
																														_snake_case													=			'\n'.join(lowerCAmelCase_					)
																														if special_strings is not None:
																																			for string in special_strings:
																																								_snake_case													=			diff.replace(lowerCAmelCase_			,  ''					)
																														self.assertEqual(lowerCAmelCase_			,  ''					)




					def 		lowerCAmelCase     (						self					)						->							List[str]:
										self.one_complete_example('complete_nlp_example.py'			,  lowerCAmelCase_					)
										self.one_complete_example('complete_nlp_example.py'			,  lowerCAmelCase_					)






					def 		lowerCAmelCase     (						self					)						->							str:
										_snake_case													=			os.path.abspath(os.path.join('examples'			,  'cv_example.py'					)					)
										_snake_case													=			[
										    ' ' * 16 + '{\n\n',
										    ' ' * 20 + '"accuracy": eval_metric["accuracy"],\n\n',
										    ' ' * 20 + '"f1": eval_metric["f1"],\n\n',
										    ' ' * 20 + '"train_loss": total_loss.item() / len(train_dataloader),\n\n',
										    ' ' * 20 + '"epoch": epoch,\n\n',
										    ' ' * 16 + '},\n\n',
										    ' ' * 16 + 'step=epoch,\n',
										    ' ' * 12,
										    ' ' * 8 + 'for step, batch in enumerate(active_dataloader):\n',
										]
										self.one_complete_example('complete_cv_example.py'			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)
										self.one_complete_example('complete_cv_example.py'			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)





@mock.patch.dict(os.environ							,					{'''TESTING_MOCKED_DATALOADERS''': '''1'''}   )
class      UpperCamelCase_						(					_lowerCamelCase   ):
					lowerCAmelCase_							=		False




					@classmethod
					def 		lowerCAmelCase     (						cls					)						->							Optional[Any]:
										super().setUpClass()
										_snake_case													=			tempfile.mkdtemp()
										_snake_case													=			os.path.join(cls._tmpdir			,  'default_config.yml'					)

										write_basic_config(save_location=cls.configPath					)
										_snake_case													=			['accelerate', 'launch', '--config_file', cls.configPath]




					@classmethod
					def 		lowerCAmelCase     (						cls					)						->							Optional[int]:
										super().tearDownClass()
										shutil.rmtree(cls._tmpdir					)




					def 		lowerCAmelCase     (						self					)						->							Optional[int]:
										_snake_case													=			F'''
        examples/by_feature/checkpointing.py
        --checkpointing_steps epoch
        --output_dir {self.tmpdir}
        '''.split()
										run_command(self._launch_args + testargs					)
										self.assertTrue(os.path.exists(os.path.join(self.tmpdir			,  'epoch_0'					)					)					)




					def 		lowerCAmelCase     (						self					)						->							int:
										_snake_case													=			F'''
        examples/by_feature/checkpointing.py
        --checkpointing_steps 1
        --output_dir {self.tmpdir}
        '''.split()
										_snake_case													=			run_command(self._launch_args + testargs					)
										self.assertTrue(os.path.exists(os.path.join(self.tmpdir			,  'step_2'					)					)					)




					def 		lowerCAmelCase     (						self					)						->							Union[str, Any]:
										_snake_case													=			F'''
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir			,  "epoch_0"					)}
        '''.split()
										_snake_case													=			run_command(self._launch_args + testargs			,  return_stdout=lowerCAmelCase_					)
										self.assertNotIn('epoch 0:'			,  lowerCAmelCase_					)
										self.assertIn('epoch 1:'			,  lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							Dict:
										_snake_case													=			F'''
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir			,  "step_2"					)}
        '''.split()
										_snake_case													=			run_command(self._launch_args + testargs			,  return_stdout=lowerCAmelCase_					)
										if torch.cuda.is_available():
															_snake_case													=			torch.cuda.device_count()
										else:
															_snake_case													=			1
										if num_processes > 1:
															self.assertNotIn('epoch 0:'			,  lowerCAmelCase_					)
															self.assertIn('epoch 1:'			,  lowerCAmelCase_					)
										else:
															self.assertIn('epoch 0:'			,  lowerCAmelCase_					)
															self.assertIn('epoch 1:'			,  lowerCAmelCase_					)




					@slow
					def 		lowerCAmelCase     (						self					)						->							Optional[int]:
										_snake_case													=			'\n        examples/by_feature/cross_validation.py\n        --num_folds 2\n        '.split()
										with mock.patch.dict(os.environ			,  {'TESTING_MOCKED_DATALOADERS': '0'}					):
															_snake_case													=			run_command(self._launch_args + testargs			,  return_stdout=lowerCAmelCase_					)
															_snake_case													=			re.findall('({.+})'			,  lowerCAmelCase_					)
															_snake_case													=			[r for r in results if 'accuracy' in r][-1]
															_snake_case													=			ast.literal_eval(lowerCAmelCase_					)
															self.assertGreaterEqual(results['accuracy']			,  0.75					)




					def 		lowerCAmelCase     (						self					)						->							Optional[int]:
										_snake_case													=			['examples/by_feature/multi_process_metrics.py']
										run_command(self._launch_args + testargs					)




					@require_trackers
					@mock.patch.dict(os.environ			,  {'WANDB_MODE': 'offline'}					)
					def 		lowerCAmelCase     (						self					)						->							Optional[Any]:
										with tempfile.TemporaryDirectory() as tmpdir:
															_snake_case													=			F'''
            examples/by_feature/tracking.py
            --with_tracking
            --project_dir {tmpdir}
            '''.split()
															run_command(self._launch_args + testargs					)
															self.assertTrue(os.path.exists(os.path.join(lowerCAmelCase_			,  'tracking'					)					)					)




					def 		lowerCAmelCase     (						self					)						->							int:
										_snake_case													=			['examples/by_feature/gradient_accumulation.py']
										run_command(self._launch_args + testargs					)






					def 		lowerCAmelCase     (						self					)						->							List[str]:
										_snake_case													=			['examples/by_feature/local_sgd.py']
										run_command(self._launch_args + testargs					)










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_ 							=       [
    {"""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   lowerCamelCase__       (       UpperCamelCase__     :    Dict=True  )     ->	Dict:



					'''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=_lowerCamelCase   )   )
class      UpperCamelCase_						(					_lowerCamelCase   ):
					lowerCAmelCase_							=		None
					lowerCAmelCase_							=		None




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							Union[str, Any]:
										with TemporaryDirectory() as tmp_dir:
															_snake_case													=			dataset_module_factory(lowerCAmelCase_			,  cache_dir=lowerCAmelCase_					)

															_snake_case													=			import_main_class(dataset_module.module_path			,  dataset=lowerCAmelCase_					)

															_snake_case													=			builder_cls(
															    cache_dir=lowerCAmelCase_			,  config_name=lowerCAmelCase_			,  hash=dataset_module.hash			,  )

															_snake_case													=			'/'.join(
															    [
															        HF_GCP_BASE_URL,
															        builder_instance._relative_data_dir(with_hash=lowerCAmelCase_					).replace(os.sep			,  '/'					),
															        config.DATASET_INFO_FILENAME,
															    ]					)
															_snake_case													=			cached_path(lowerCAmelCase_			,  cache_dir=lowerCAmelCase_					)
															self.assertTrue(os.path.exists(lowerCAmelCase_					)					)





@pytest.mark.integration
def   lowerCamelCase__       (       UpperCamelCase__     :    Any  )     ->	Tuple:



					'''simple docstring'''




					_snake_case													=			tmp_path_factory.mktemp('test_hf_gcp'  ) / 'test_wikipedia_simple'
					_snake_case													=			dataset_module_factory('wikipedia'    ,      cache_dir=UpperCamelCase__  )
					_snake_case													=			import_main_class(dataset_module.module_path  )
					_snake_case													=			builder_cls(
					    cache_dir=UpperCamelCase__    ,      config_name='20220301.frr'    ,      hash=dataset_module.hash    ,      )
					# use the HF cloud storage, not the original download_and_prepare that uses apache-beam
					_snake_case													=			None
					builder_instance.download_and_prepare()
					_snake_case													=			builder_instance.as_dataset()
					assert ds





@pytest.mark.integration
def   lowerCamelCase__       (       UpperCamelCase__     :    Union[str, Any]  )     ->	Union[str, Any]:



					'''simple docstring'''




					_snake_case													=			dataset_module_factory('wikipedia'    ,      cache_dir=UpperCamelCase__  )
					_snake_case													=			import_main_class(dataset_module.module_path    ,      dataset=UpperCamelCase__  )
					_snake_case													=			builder_cls(
					    cache_dir=UpperCamelCase__    ,      config_name='20220301.frr'    ,      hash=dataset_module.hash    ,      )
					_snake_case													=			builder_instance.as_streaming_dataset()
					assert ds
					assert isinstance(UpperCamelCase__    ,      UpperCamelCase__  )
					assert "train" in ds
					assert isinstance(ds['train']    ,      UpperCamelCase__  )
					assert next(iter(ds['train']  )  )



