"""simple docstring"""




import warnings
from typing import List
from unittest.mock import Mock

import torch
from torch.utils.data import DataLoader, IterableDataset, TensorDataset

from accelerate.accelerator import Accelerator
from accelerate.utils.dataclasses import DistributedType


class      lowerCamelCase_					(       lowerCamelCase__			):






    '''simple docstring'''






    def __init__(     self			,			__lowercase)	->     Optional[int]:
        __UpperCamelCase					:int        =		data







    def __iter__(     self)	->     List[str]:
        for element in self.data:
            yield element



def lowerCamelCase  (       SCREAMING_SNAKE_CASE=True     ):
    '''simple docstring'''







    __UpperCamelCase					:int        =		Accelerator(even_batches=_A     )
    assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
    return accelerator



def lowerCamelCase  (       SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE = False     ):
    '''simple docstring'''







    if iterable:
        __UpperCamelCase					:Optional[int]        =		DummyIterableDataset(torch.as_tensor(range(_A     )     )     )
    else:
        __UpperCamelCase					:Tuple        =		TensorDataset(torch.as_tensor(range(_A     )     )     )

    __UpperCamelCase					:List[Any]        =		DataLoader(_A  ,      batch_size=_A     )
    __UpperCamelCase					:int        =		accelerator.prepare(_A     )

    return dl



def lowerCamelCase  (       SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE  ,      ):
    '''simple docstring'''







    __UpperCamelCase					:Tuple        =		create_dataloader(accelerator=_A  ,      dataset_size=_A  ,      batch_size=_A     )

    __UpperCamelCase					:List[str]        =		[len(batch[0]     ) for batch in dl]

    if accelerator.process_index == 0:
        assert batch_sizes == process_0_expected_batch_sizes
    elif accelerator.process_index == 1:
        assert batch_sizes == process_1_expected_batch_sizes



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:Tuple        =		create_accelerator()

    # without padding, we would expect a different number of batches
    verify_dataloader_batch_sizes(
        _A  ,      dataset_size=3  ,      batch_size=1  ,      process_0_expected_batch_sizes=[1, 1]  ,      process_1_expected_batch_sizes=[1, 1]  ,      )

    # without padding, we would expect the same number of batches, but different sizes
    verify_dataloader_batch_sizes(
        _A  ,      dataset_size=7  ,      batch_size=2  ,      process_0_expected_batch_sizes=[2, 2]  ,      process_1_expected_batch_sizes=[2, 2]  ,      )



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:List[str]        =		create_accelerator(even_batches=_A     )

    verify_dataloader_batch_sizes(
        _A  ,      dataset_size=3  ,      batch_size=1  ,      process_0_expected_batch_sizes=[1, 1]  ,      process_1_expected_batch_sizes=[1]  ,      )

    verify_dataloader_batch_sizes(
        _A  ,      dataset_size=7  ,      batch_size=2  ,      process_0_expected_batch_sizes=[2, 2]  ,      process_1_expected_batch_sizes=[2, 1]  ,      )



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:List[Any]        =		create_accelerator(even_batches=_A     )

    __UpperCamelCase					:str        =		torch.nn.Linear(1  ,      1     )
    __UpperCamelCase					:Any        =		accelerator.prepare(_A     )

    __UpperCamelCase					:List[Any]        =		create_dataloader(_A  ,      dataset_size=3  ,      batch_size=1     )

    __UpperCamelCase					:List[Any]        =		[]
    with accelerator.join_uneven_inputs([ddp_model]     ):
        for batch_idx, batch in enumerate(_A     ):
            __UpperCamelCase					:List[Any]        =		ddp_model(batch[0].float()     )
            __UpperCamelCase					:Optional[int]        =		output.sum()
            loss.backward()
            batch_idxs.append(_A     )

    accelerator.wait_for_everyone()

    if accelerator.process_index == 0:
        assert batch_idxs == [0, 1]
    elif accelerator.process_index == 1:
        assert batch_idxs == [0]



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    with warnings.catch_warnings(record=_A     ) as w:
        with accelerator.join_uneven_inputs([Mock()]     ):
            pass

        assert issubclass(w[-1].category  ,      _A     )
        assert "only supported for multi-GPU" in str(w[-1].message     )



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:Optional[Any]        =		True
    __UpperCamelCase					:List[str]        =		False
    __UpperCamelCase					:Tuple        =		create_accelerator(even_batches=_A     )
    __UpperCamelCase					:Optional[Any]        =		torch.nn.Linear(1  ,      1     )
    __UpperCamelCase					:List[Any]        =		accelerator.prepare(_A     )
    __UpperCamelCase					:Optional[int]        =		create_dataloader(_A  ,      dataset_size=3  ,      batch_size=1     )
    __UpperCamelCase					:Dict        =		create_dataloader(_A  ,      dataset_size=3  ,      batch_size=1     )

    with accelerator.join_uneven_inputs([ddp_model]  ,      even_batches=_A     ):
        __UpperCamelCase					:str        =		train_dl.batch_sampler.even_batches
        __UpperCamelCase					:Optional[Any]        =		valid_dl.batch_sampler.even_batches

    assert train_dl_overridden_value == overridden_even_batches
    assert valid_dl_overridden_value == overridden_even_batches
    assert train_dl.batch_sampler.even_batches == default_even_batches
    assert valid_dl.batch_sampler.even_batches == default_even_batches



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:int        =		True
    __UpperCamelCase					:int        =		False
    __UpperCamelCase					:Union[str, Any]        =		create_accelerator(even_batches=_A     )
    __UpperCamelCase					:List[Any]        =		torch.nn.Linear(1  ,      1     )
    __UpperCamelCase					:Optional[int]        =		accelerator.prepare(_A     )
    create_dataloader(_A  ,      dataset_size=3  ,      batch_size=1  ,      iterable=_A     )
    __UpperCamelCase					:List[str]        =		create_dataloader(_A  ,      dataset_size=3  ,      batch_size=1     )

    with warnings.catch_warnings():
        warnings.filterwarnings('''ignore'''     )
        try:
            with accelerator.join_uneven_inputs([ddp_model]  ,      even_batches=_A     ):
                __UpperCamelCase					:Any        =		batch_dl.batch_sampler.even_batches
        except AttributeError:
            # ensure attribute error is not raised when processing iterable dl
            raise AssertionError

    assert batch_dl_overridden_value == overridden_even_batches
    assert batch_dl.batch_sampler.even_batches == default_even_batches



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:Dict        =		create_accelerator()
    __UpperCamelCase					:List[Any]        =		torch.nn.Linear(1  ,      1     )
    __UpperCamelCase					:Tuple        =		accelerator.prepare(_A     )
    create_dataloader(_A  ,      dataset_size=3  ,      batch_size=1  ,      iterable=_A     )

    with warnings.catch_warnings(record=_A     ) as w:
        with accelerator.join_uneven_inputs([ddp_model]  ,      even_batches=_A     ):
            pass

        assert issubclass(w[-1].category  ,      _A     )
        assert "only supported for map-style datasets" in str(w[-1].message     )



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:Optional[Any]        =		create_accelerator()

    accelerator.print('''Test that even_batches variable ensures uniform batches across processes'''     )
    test_default_ensures_even_batch_sizes()

    accelerator.print('''Run tests with even_batches disabled'''     )
    test_can_disable_even_batches()

    accelerator.print('''Test joining uneven inputs'''     )
    test_can_join_uneven_inputs()

    accelerator.print('''Test overriding even_batches when joining uneven inputs'''     )
    test_join_can_override_even_batches()

    accelerator.print('''Test overriding even_batches for mixed dataloader types'''     )
    test_join_can_override_for_mixed_type_dataloaders()

    accelerator.print('''Test overriding even_batches raises a warning for iterable dataloaders'''     )
    test_join_raises_warning_for_iterable_when_overriding_even_batches()

    accelerator.print('''Test join with non DDP distributed raises warning'''     )
    __UpperCamelCase					:Any        =		accelerator.state.distributed_type
    __UpperCamelCase					:Dict        =		DistributedType.FSDP
    test_join_raises_warning_for_non_ddp_distributed(_A     )
    __UpperCamelCase					:str        =		original_state


if __name__ == "__main__":
    main()






import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


sys.path.append(str(Path(__file__).parent.parent.parent.parent / '''utils'''))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402


class      lowerCamelCase_					(       unittest.TestCase			):






    '''simple docstring'''






    def       UpperCamelCase__       (     self)	->     int:
        __UpperCamelCase					:str        =		0







    def       UpperCamelCase__       (     self)	->     Optional[Any]:
        __UpperCamelCase					:Dict        =		AutoImageProcessor.from_pretrained('''openai/clip-vit-base-patch32''')
        self.assertIsInstance(__lowercase			,			__lowercase)







    def       UpperCamelCase__       (     self)	->     Optional[int]:
        with tempfile.TemporaryDirectory() as tmpdirname:
            __UpperCamelCase					:int        =		Path(__lowercase) / '''preprocessor_config.json'''
            __UpperCamelCase					:Dict        =		Path(__lowercase) / '''config.json'''
            json.dump(
                {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			,			open(__lowercase			,			'''w''')			,			)
            json.dump({'''model_type''': '''clip'''}			,			open(__lowercase			,			'''w'''))

            __UpperCamelCase					:Union[str, Any]        =		AutoImageProcessor.from_pretrained(__lowercase)
            self.assertIsInstance(__lowercase			,			__lowercase)







    def       UpperCamelCase__       (     self)	->     Union[str, Any]:
        # Ensure we can load the image processor from the feature extractor config
        with tempfile.TemporaryDirectory() as tmpdirname:
            __UpperCamelCase					:str        =		Path(__lowercase) / '''preprocessor_config.json'''
            __UpperCamelCase					:Union[str, Any]        =		Path(__lowercase) / '''config.json'''
            json.dump(
                {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}			,			open(__lowercase			,			'''w''')			,			)
            json.dump({'''model_type''': '''clip'''}			,			open(__lowercase			,			'''w'''))

            __UpperCamelCase					:Dict        =		AutoImageProcessor.from_pretrained(__lowercase)
            self.assertIsInstance(__lowercase			,			__lowercase)







    def       UpperCamelCase__       (     self)	->     Optional[int]:
        with tempfile.TemporaryDirectory() as tmpdirname:
            __UpperCamelCase					:int        =		CLIPConfig()

            # Create a dummy config file with image_proceesor_type
            __UpperCamelCase					:Tuple        =		Path(__lowercase) / '''preprocessor_config.json'''
            __UpperCamelCase					:Optional[Any]        =		Path(__lowercase) / '''config.json'''
            json.dump(
                {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			,			open(__lowercase			,			'''w''')			,			)
            json.dump({'''model_type''': '''clip'''}			,			open(__lowercase			,			'''w'''))

            # remove image_processor_type to make sure config.json alone is enough to load image processor locally
            __UpperCamelCase					:Optional[Any]        =		AutoImageProcessor.from_pretrained(__lowercase).to_dict()

            config_dict.pop('''image_processor_type''')
            __UpperCamelCase					:List[str]        =		CLIPImageProcessor(**__lowercase)

            # save in new folder
            model_config.save_pretrained(__lowercase)
            config.save_pretrained(__lowercase)

            __UpperCamelCase					:Dict        =		AutoImageProcessor.from_pretrained(__lowercase)

            # make sure private variable is not incorrectly saved
            __UpperCamelCase					:Union[str, Any]        =		json.loads(config.to_json_string())
            self.assertTrue('''_processor_class''' not in dict_as_saved)

        self.assertIsInstance(__lowercase			,			__lowercase)







    def       UpperCamelCase__       (     self)	->     List[str]:
        with tempfile.TemporaryDirectory() as tmpdirname:
            __UpperCamelCase					:Tuple        =		Path(__lowercase) / '''preprocessor_config.json'''
            json.dump(
                {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			,			open(__lowercase			,			'''w''')			,			)

            __UpperCamelCase					:List[str]        =		AutoImageProcessor.from_pretrained(__lowercase)
            self.assertIsInstance(__lowercase			,			__lowercase)







    def       UpperCamelCase__       (     self)	->     Optional[int]:
        with self.assertRaisesRegex(
            __lowercase			,			'''clip-base is not a local folder and is not a valid model identifier'''):
            __UpperCamelCase					:Optional[Any]        =		AutoImageProcessor.from_pretrained('''clip-base''')







    def       UpperCamelCase__       (     self)	->     List[Any]:
        with self.assertRaisesRegex(
            __lowercase			,			r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''):
            __UpperCamelCase					:str        =		AutoImageProcessor.from_pretrained(__lowercase			,			revision='''aaaaaa''')







    def       UpperCamelCase__       (     self)	->     List[str]:
        with self.assertRaisesRegex(
            __lowercase			,			'''hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'''			,			):
            __UpperCamelCase					:Optional[Any]        =		AutoImageProcessor.from_pretrained('''hf-internal-testing/config-no-model''')







    def       UpperCamelCase__       (     self)	->     str:
        # If remote code is not set, we will time out when asking whether to load the model.
        with self.assertRaises(__lowercase):
            __UpperCamelCase					:Dict        =		AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor''')
        # If remote code is disabled, we can't load this config.
        with self.assertRaises(__lowercase):
            __UpperCamelCase					:List[Any]        =		AutoImageProcessor.from_pretrained(
                '''hf-internal-testing/test_dynamic_image_processor'''			,			trust_remote_code=__lowercase)

        __UpperCamelCase					:Optional[Any]        =		AutoImageProcessor.from_pretrained(
            '''hf-internal-testing/test_dynamic_image_processor'''			,			trust_remote_code=__lowercase)
        self.assertEqual(image_processor.__class__.__name__			,			'''NewImageProcessor''')

        # Test image processor can be reloaded.
        with tempfile.TemporaryDirectory() as tmp_dir:
            image_processor.save_pretrained(__lowercase)
            __UpperCamelCase					:List[Any]        =		AutoImageProcessor.from_pretrained(__lowercase			,			trust_remote_code=__lowercase)
        self.assertEqual(reloaded_image_processor.__class__.__name__			,			'''NewImageProcessor''')







    def       UpperCamelCase__       (     self)	->     Optional[Any]:
        try:
            AutoConfig.register('''custom'''			,			__lowercase)
            AutoImageProcessor.register(__lowercase			,			__lowercase)
            # Trying to register something existing in the Transformers library will raise an error
            with self.assertRaises(__lowercase):
                AutoImageProcessor.register(__lowercase			,			__lowercase)

            with tempfile.TemporaryDirectory() as tmpdirname:
                __UpperCamelCase					:int        =		Path(__lowercase) / '''preprocessor_config.json'''
                __UpperCamelCase					:List[str]        =		Path(__lowercase) / '''config.json'''
                json.dump(
                    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}			,			open(__lowercase			,			'''w''')			,			)
                json.dump({'''model_type''': '''clip'''}			,			open(__lowercase			,			'''w'''))

                __UpperCamelCase					:int        =		CustomImageProcessor.from_pretrained(__lowercase)

            # Now that the config is registered, it can be used as any other config with the auto-API
            with tempfile.TemporaryDirectory() as tmp_dir:
                image_processor.save_pretrained(__lowercase)
                __UpperCamelCase					:int        =		AutoImageProcessor.from_pretrained(__lowercase)
                self.assertIsInstance(__lowercase			,			__lowercase)

        finally:
            if "custom" in CONFIG_MAPPING._extra_content:
                del CONFIG_MAPPING._extra_content["custom"]
            if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]







    def       UpperCamelCase__       (     self)	->     List[Any]:


        class      lowerCamelCase_					(       UpperCAmelCase_			):






            '''simple docstring'''






            a__     :						List[str]      =				True

        try:
            AutoConfig.register('''custom'''			,			__lowercase)
            AutoImageProcessor.register(__lowercase			,			__lowercase)
            # If remote code is not set, the default is to use local
            __UpperCamelCase					:str        =		AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor''')
            self.assertEqual(image_processor.__class__.__name__			,			'''NewImageProcessor''')
            self.assertTrue(image_processor.is_local)

            # If remote code is disabled, we load the local one.
            __UpperCamelCase					:Optional[Any]        =		AutoImageProcessor.from_pretrained(
                '''hf-internal-testing/test_dynamic_image_processor'''			,			trust_remote_code=__lowercase)
            self.assertEqual(image_processor.__class__.__name__			,			'''NewImageProcessor''')
            self.assertTrue(image_processor.is_local)

            # If remote is enabled, we load from the Hub
            __UpperCamelCase					:List[str]        =		AutoImageProcessor.from_pretrained(
                '''hf-internal-testing/test_dynamic_image_processor'''			,			trust_remote_code=__lowercase)
            self.assertEqual(image_processor.__class__.__name__			,			'''NewImageProcessor''')
            self.assertTrue(not hasattr(__lowercase			,			'''is_local'''))

        finally:
            if "custom" in CONFIG_MAPPING._extra_content:
                del CONFIG_MAPPING._extra_content["custom"]
            if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]





