




'''simple docstring'''



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 							_a				(			unittest.TestCase ):





       '''simple docstring'''






       def        UpperCamelCase_							(      self,							A,							A,							A = None,							A = None  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						None
             SCREAMING_SNAKE_CASE   :      str    					=						os.path.abspath(os.path.join('examples',							'by_feature'  )  )
             SCREAMING_SNAKE_CASE   :      str    					=						os.path.abspath('examples'  )
             for item in os.listdir(A  ):
                   if item not in EXCLUDE_EXAMPLES:
                         SCREAMING_SNAKE_CASE   :      Optional[int]    					=						os.path.join(A,							A  )
                         if os.path.isfile(A  ) and ".py" in item_path:
                               with self.subTest(
                                   tested_script=A,							feature_script=A,							tested_section='main()' if parser_only else 'training_function()',							):
                                     SCREAMING_SNAKE_CASE   :      str    					=						compare_against_test(
                                         os.path.join(A,							A  ),							A,							A,							A  )
                                     SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						'\n'.join(A  )
                                     if special_strings is not None:
                                           for string in special_strings:
                                                 SCREAMING_SNAKE_CASE   :      int    					=						diff.replace(A,							''  )
                                     self.assertEqual(A,							''  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             self.one_complete_example('complete_nlp_example.py',							A  )
             self.one_complete_example('complete_nlp_example.py',							A  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Dict    					=						os.path.abspath(os.path.join('examples',							'cv_example.py'  )  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						[
                 ' ' * 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',							A,							A,							A  )
             self.one_complete_example('complete_cv_example.py',							A,							A,							A  )






@mock.patch.dict(os.environ ,      {'''TESTING_MOCKED_DATALOADERS''': '''1'''} )
class 							_a				(			SCREAMING_SNAKE_CASE ):





       '''simple docstring'''






       A						:				List[Any]        =						False







       @classmethod
       def        UpperCamelCase_							(      cls  ):







             '''simple docstring'''







             super().setUpClass()
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						tempfile.mkdtemp()
             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						os.path.join(cls._tmpdir,							'default_config.yml'  )

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







       @classmethod
       def        UpperCamelCase_							(      cls  ):







             '''simple docstring'''







             super().tearDownClass()
             shutil.rmtree(cls._tmpdir  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						F"\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps epoch\n        --output_dir {self.tmpdir}\n        ".split()
             run_command(self._launch_args + testargs  )
             self.assertTrue(os.path.exists(os.path.join(self.tmpdir,							'epoch_0'  )  )  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Tuple    					=						F"\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps 1\n        --output_dir {self.tmpdir}\n        ".split()
             SCREAMING_SNAKE_CASE   :      Tuple    					=						run_command(self._launch_args + testargs  )
             self.assertTrue(os.path.exists(os.path.join(self.tmpdir,							'step_2'  )  )  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      str    					=						F"\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir,							'epoch_0'  )}\n        ".split()
             SCREAMING_SNAKE_CASE   :      List[str]    					=						run_command(self._launch_args + testargs,							return_stdout=A  )
             self.assertNotIn('epoch 0:',							A  )
             self.assertIn('epoch 1:',							A  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						F"\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir,							'step_2'  )}\n        ".split()
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						run_command(self._launch_args + testargs,							return_stdout=A  )
             if torch.cuda.is_available():
                   SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						torch.cuda.device_count()
             else:
                   SCREAMING_SNAKE_CASE   :      Any    					=						1
             if num_processes > 1:
                   self.assertNotIn('epoch 0:',							A  )
                   self.assertIn('epoch 1:',							A  )
             else:
                   self.assertIn('epoch 0:',							A  )
                   self.assertIn('epoch 1:',							A  )







       @slow
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						'\n        examples/by_feature/cross_validation.py\n        --num_folds 2\n        '.split()
             with mock.patch.dict(os.environ,							{'TESTING_MOCKED_DATALOADERS': '0'}  ):
                   SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						run_command(self._launch_args + testargs,							return_stdout=A  )
                   SCREAMING_SNAKE_CASE   :      Optional[int]    					=						re.findall('({.+})',							A  )
                   SCREAMING_SNAKE_CASE   :      Dict    					=						[r for r in results if 'accuracy' in r][-1]
                   SCREAMING_SNAKE_CASE   :      Optional[int]    					=						ast.literal_eval(A  )
                   self.assertGreaterEqual(results['accuracy'],							0.75  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						['examples/by_feature/multi_process_metrics.py']
             run_command(self._launch_args + testargs  )







       @require_trackers
       @mock.patch.dict(os.environ,							{'WANDB_MODE': 'offline'}  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             with tempfile.TemporaryDirectory() as tmpdir:
                   SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						F"\n            examples/by_feature/tracking.py\n            --with_tracking\n            --project_dir {tmpdir}\n            ".split()
                   run_command(self._launch_args + testargs  )
                   self.assertTrue(os.path.exists(os.path.join(A,							'tracking'  )  )  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      List[Any]    					=						['examples/by_feature/gradient_accumulation.py']
             run_command(self._launch_args + testargs  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      List[Any]    					=						['examples/by_feature/local_sgd.py']
             run_command(self._launch_args + testargs  )




import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput




class   __lowercase						(						__snake_case       ):






	def __init__(self				: List[Any]      ,	snake_case				: Optional[int]      ,	snake_case				: str      ,	snake_case				: str    ) ->			Dict:
							_lowercase       :	Tuple           =   dataset
							_lowercase       :	List[str]           =   process
							_lowercase       :	Any           =   params






	def __len__(self				: Optional[Any]    ) ->			Any:
							return len(self.dataset    )






	def __getitem__(self				: int      ,	snake_case				: Any    ) ->			int:
							_lowercase       :	Optional[Any]           =   self.dataset[i]
							_lowercase       :	Union[str, Any]           =   self.process(snake_case      ,	**self.params    )
							return processed




class   __lowercase						(						__snake_case       ):






	def __init__(self				: int      ,	snake_case				: List[Any]      ,	snake_case				: Tuple      ,	snake_case				: Tuple      ,	snake_case				: List[str]=None    ) ->			Optional[int]:
							_lowercase       :	List[str]           =   loader
							_lowercase       :	Optional[Any]           =   infer
							_lowercase       :	List[Any]           =   params
							if loader_batch_size == 1:
													# Let's spare some time by deactivating altogether
													_lowercase       :	str           =   None
							_lowercase       :	Optional[int]           =   loader_batch_size

							# Internal bookkeeping
							_lowercase       :	Dict           =   None
							_lowercase       :	Union[str, Any]           =   None






	def __len__(self				: Any    ) ->			Any:
							return len(self.loader    )






	def __iter__(self				: int    ) ->			Optional[Any]:
							_lowercase       :	List[str]           =   iter(self.loader    )
							return self






	def       _a(self				: int    ) ->			List[str]:
							if isinstance(self._loader_batch_data      ,	torch.Tensor    ):
													# Batch data is simple tensor, just fetch the slice
													_lowercase       :	Any           =   self._loader_batch_data[self._loader_batch_index]
							else:
													# Batch data is assumed to be BaseModelOutput (or dict)
													_lowercase       :	List[str]           =   {}
													for k, element in self._loader_batch_data.items():
																			if isinstance(snake_case      ,	snake_case    ):
																									# Convert ModelOutput to tuple first
																									_lowercase       :	Dict           =   element.to_tuple()
																									if isinstance(element[0]      ,	torch.Tensor    ):
																															_lowercase       :	Optional[int]           =   tuple(el[self._loader_batch_index].unsqueeze(0    ) for el in element    )
																									elif isinstance(element[0]      ,	np.ndarray    ):
																															_lowercase       :	Optional[Any]           =   tuple(np.expand_dims(el[self._loader_batch_index]      ,	0    ) for el in element    )
																									continue
																			if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(snake_case      ,	snake_case    ):
																									# Those are stored as lists of tensors so need specific unbatching.
																									if isinstance(element[0]      ,	torch.Tensor    ):
																															_lowercase       :	Optional[int]           =   tuple(el[self._loader_batch_index].unsqueeze(0    ) for el in element    )
																									elif isinstance(element[0]      ,	np.ndarray    ):
																															_lowercase       :	List[str]           =   tuple(np.expand_dims(el[self._loader_batch_index]      ,	0    ) for el in element    )
																									continue
																			if element is None:
																									# This can happen for optional data that get passed around
																									_lowercase       :	Optional[Any]           =   None
																			elif isinstance(element[self._loader_batch_index]      ,	torch.Tensor    ):
																									# Take correct batch data, but make it looked like batch_size=1
																									# For compatibility with other methods within transformers

																									_lowercase       :	int           =   element[self._loader_batch_index].unsqueeze(0    )
																			elif isinstance(element[self._loader_batch_index]      ,	np.ndarray    ):
																									# Take correct batch data, but make it looked like batch_size=1
																									# For compatibility with other methods within transformers
																									_lowercase       :	Tuple           =   np.expand_dims(element[self._loader_batch_index]      ,	0    )
																			else:
																									# This is typically a list, so no need to `unsqueeze`.
																									_lowercase       :	Union[str, Any]           =   element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
													_lowercase       :	int           =   self._loader_batch_data.__class__(snake_case    )
							self._loader_batch_index += 1
							return result






	def       _a(self				: Any    ) ->			str:
							if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
													# We are currently unrolling a batch so we just need to return
													# the current item within a batch
													return self.loader_batch_item()

							# We're out of items within a batch
							_lowercase       :	str           =   next(self.iterator    )
							_lowercase       :	Any           =   self.infer(snake_case      ,	**self.params    )
							# We now have a batch of "inferred things".
							if self.loader_batch_size is not None:
													# Try to infer the size of the batch
													if isinstance(snake_case      ,	torch.Tensor    ):
																			_lowercase       :	int           =   processed
													else:
																			_lowercase       :	List[str]           =   list(processed.keys()    )[0]
																			_lowercase       :	List[str]           =   processed[key]
													if isinstance(snake_case      ,	snake_case    ):
																			_lowercase       :	int           =   len(snake_case    )
													else:
																			_lowercase       :	Optional[int]           =   first_tensor.shape[0]
													if 0 < observed_batch_size < self.loader_batch_size:
																			# could be last batch so we can't unroll as many
																			# elements.
																			_lowercase       :	Union[str, Any]           =   observed_batch_size
													# Setting internal index to unwrap the batch
													_lowercase       :	Any           =   processed
													_lowercase       :	Dict           =   0
													return self.loader_batch_item()
							else:
													# We're not unrolling batches
													return processed




class   __lowercase						(						__snake_case       ):






	def __init__(self				: Union[str, Any]      ,	snake_case				: str      ,	snake_case				: Dict      ,	snake_case				: Optional[Any]      ,	snake_case				: List[str]=None    ) ->			List[str]:
							super().__init__(snake_case      ,	snake_case      ,	snake_case    )






	def __iter__(self				: List[Any]    ) ->			str:
							_lowercase       :	int           =   iter(self.loader    )
							_lowercase       :	Optional[int]           =   None
							return self






	def       _a(self				: Optional[Any]    ) ->			Optional[int]:
							if self.subiterator is None:
													_lowercase       :	List[str]           =   self.infer(next(self.iterator    )      ,	**self.params    )
							try:
													# Try to return next item
													_lowercase       :	str           =   next(self.subiterator    )
							except StopIteration:
													# When a preprocess iterator ends, we can start lookig at the next item
													# ChunkIterator will keep feeding until ALL elements of iterator
													# all have created their subiterator and have been iterating against.
													#
													# Another way to look at it, is we're basically flattening lists of lists
													# into a single list, but with generators
													_lowercase       :	Any           =   self.infer(next(self.iterator    )      ,	**self.params    )
													_lowercase       :	Optional[int]           =   next(self.subiterator    )
							return processed




class   __lowercase						(						__snake_case       ):






	def __iter__(self				: List[str]    ) ->			List[str]:
							_lowercase       :	Any           =   iter(self.loader    )
							return self






	def       _a(self				: List[Any]    ) ->			Union[str, Any]:
							# Extremely similar to PipelineIterator in its unpacking mechanism
							# BUT, we have an extra required item which is the presence of `is_last`
							# That is because everything is flattened by `PipelineChunkIterator` we
							# need to keep track of how to regroup here in the original `process`
							# boundaries so that `process` and `postprocess` see the same data.

							# This iterator accumulates items (possibly while unbatching) until it
							# its a `is_last` and then just passes it on to the caller.
							_lowercase       :	Dict           =   False
							_lowercase       :	Optional[int]           =   []
							if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
													while self._loader_batch_index < self.loader_batch_size:
																			_lowercase       :	Tuple           =   self.loader_batch_item()
																			_lowercase       :	List[Any]           =   item.pop("is_last"    )
																			accumulator.append(snake_case    )
																			if is_last:
																									return accumulator

							while not is_last:
													_lowercase       :	List[Any]           =   self.infer(next(self.iterator    )      ,	**self.params    )
													if self.loader_batch_size is not None:
																			if isinstance(snake_case      ,	torch.Tensor    ):
																									_lowercase       :	Optional[Any]           =   processed
																			else:
																									_lowercase       :	Tuple           =   list(processed.keys()    )[0]
																									_lowercase       :	int           =   processed[key]
																			if isinstance(snake_case      ,	snake_case    ):
																									_lowercase       :	str           =   len(snake_case    )
																			else:
																									_lowercase       :	List[str]           =   first_tensor.shape[0]
																			if 0 < observed_batch_size < self.loader_batch_size:
																									# could be last batch so we can't unroll as many
																									# elements.
																									_lowercase       :	Tuple           =   observed_batch_size
																			_lowercase       :	Any           =   processed
																			_lowercase       :	Dict           =   0
																			while self._loader_batch_index < self.loader_batch_size:
																									_lowercase       :	str           =   self.loader_batch_item()
																									_lowercase       :	int           =   item.pop("is_last"    )
																									accumulator.append(snake_case    )
																									if is_last:
																															return accumulator
													else:
																			_lowercase       :	str           =   processed
																			_lowercase       :	int           =   item.pop("is_last"    )
																			accumulator.append(snake_case    )
							return accumulator




class   __lowercase						(						__snake_case       ):






	def __init__(self				: int      ,	snake_case				: Dataset      ,	snake_case				: str    ) ->			List[Any]:
							_lowercase       :	Optional[Any]           =   dataset
							_lowercase       :	Any           =   key






	def __len__(self				: Any    ) ->			Union[str, Any]:
							return len(self.dataset    )






	def __getitem__(self				: int      ,	snake_case				: Any    ) ->			Any:
							return self.dataset[i][self.key]




class   __lowercase						(						__snake_case       ):






	def __init__(self				: int      ,	snake_case				: Dataset      ,	snake_case				: str      ,	snake_case				: str    ) ->			Dict:
							_lowercase       :	int           =   dataset
							_lowercase       :	Optional[Any]           =   keya
							_lowercase       :	Tuple           =   keya






	def __len__(self				: List[str]    ) ->			Union[str, Any]:
							return len(self.dataset    )






	def __getitem__(self				: Optional[Any]      ,	snake_case				: Dict    ) ->			int:
							return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}




