




import os
import unittest
from tempfile import TemporaryDirectory

import torch
import torch.nn as nn

from accelerate.utils import (
    OffloadedWeightsLoader,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
)

class     __a(	nn.Module  ):






    """simple docstring"""





    def __init__( self    )					->						List[str]:
         super().__init__()
         UpperCAmelCase_     :				Any  =     nn.Linear(3  ,4    )
         UpperCAmelCase_     :				int  =     nn.BatchNormad(4    )
         UpperCAmelCase_     :				List[Any]  =     nn.Linear(4  ,5    )





    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE    )					->						str:
         return self.lineara(self.batchnorm(self.lineara(_SCREAMING_SNAKE_CASE    )    )    )


class     __a(	unittest.TestCase  ):






    """simple docstring"""





    def 				a__ ( self    )					->						Tuple:
         UpperCAmelCase_     :				int  =     ModelForTest()
         with TemporaryDirectory() as tmp_dir:
              offload_state_dict(_SCREAMING_SNAKE_CASE  ,model.state_dict()    )
              UpperCAmelCase_     :				Optional[int]  =     os.path.join(_SCREAMING_SNAKE_CASE  ,'''index.json'''    )
              self.assertTrue(os.path.isfile(_SCREAMING_SNAKE_CASE    )    )
              # TODO: add tests on what is inside the index

              for key in ["linear1.weight", "linear1.bias", "linear2.weight", "linear2.bias"]:
                   UpperCAmelCase_     :				Dict  =     os.path.join(_SCREAMING_SNAKE_CASE  ,f'''{key}.dat'''    )
                   self.assertTrue(os.path.isfile(_SCREAMING_SNAKE_CASE    )    )
                   # TODO: add tests on the fact weights are properly loaded





    def 				a__ ( self    )					->						int:
         UpperCAmelCase_     :				List[str]  =     [torch.floataa, torch.floataa, torch.bfloataa]

         for dtype in dtypes:
              UpperCAmelCase_     :				int  =     torch.randn(2  ,3  ,dtype=_SCREAMING_SNAKE_CASE    )
              with TemporaryDirectory() as tmp_dir:
                   UpperCAmelCase_     :				Any  =     offload_weight(_SCREAMING_SNAKE_CASE  ,'''weight'''  ,_SCREAMING_SNAKE_CASE  ,{}    )
                   UpperCAmelCase_     :				List[Any]  =     os.path.join(_SCREAMING_SNAKE_CASE  ,'''weight.dat'''    )
                   self.assertTrue(os.path.isfile(_SCREAMING_SNAKE_CASE    )    )
                   self.assertDictEqual(_SCREAMING_SNAKE_CASE  ,{'''weight''': {'''shape''': [2, 3], '''dtype''': str(_SCREAMING_SNAKE_CASE    ).split('''.'''    )[1]}}    )

                   UpperCAmelCase_     :				Tuple  =     load_offloaded_weight(_SCREAMING_SNAKE_CASE  ,index['''weight''']    )
                   self.assertTrue(torch.equal(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )    )





    def 				a__ ( self    )					->						List[Any]:
         UpperCAmelCase_     :				int  =     ModelForTest()
         UpperCAmelCase_     :				int  =     model.state_dict()
         UpperCAmelCase_     :				int  =     {k: v for k, v in state_dict.items() if '''linear2''' not in k}
         UpperCAmelCase_     :				str  =     {k: v for k, v in state_dict.items() if '''linear2''' in k}

         with TemporaryDirectory() as tmp_dir:
              offload_state_dict(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )
              UpperCAmelCase_     :				Dict  =     OffloadedWeightsLoader(state_dict=_SCREAMING_SNAKE_CASE  ,save_folder=_SCREAMING_SNAKE_CASE    )

              # Every key is there with the right value
              self.assertEqual(sorted(_SCREAMING_SNAKE_CASE    )  ,sorted(state_dict.keys()    )    )
              for key, param in state_dict.items():
                   self.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE  ,weight_map[key]    )    )

         UpperCAmelCase_     :				List[Any]  =     {k: v for k, v in state_dict.items() if '''weight''' in k}
         UpperCAmelCase_     :				Tuple  =     {k: v for k, v in state_dict.items() if '''weight''' not in k}

         with TemporaryDirectory() as tmp_dir:
              offload_state_dict(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )
              UpperCAmelCase_     :				str  =     OffloadedWeightsLoader(state_dict=_SCREAMING_SNAKE_CASE  ,save_folder=_SCREAMING_SNAKE_CASE    )

              # Every key is there with the right value
              self.assertEqual(sorted(_SCREAMING_SNAKE_CASE    )  ,sorted(state_dict.keys()    )    )
              for key, param in state_dict.items():
                   self.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE  ,weight_map[key]    )    )

         with TemporaryDirectory() as tmp_dir:
              offload_state_dict(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )
              # Duplicates are removed
              UpperCAmelCase_     :				Optional[int]  =     OffloadedWeightsLoader(state_dict=_SCREAMING_SNAKE_CASE  ,save_folder=_SCREAMING_SNAKE_CASE    )

              # Every key is there with the right value
              self.assertEqual(sorted(_SCREAMING_SNAKE_CASE    )  ,sorted(state_dict.keys()    )    )
              for key, param in state_dict.items():
                   self.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE  ,weight_map[key]    )    )





    def 				a__ ( self    )					->						List[str]:
         UpperCAmelCase_     :				int  =     {'''a.1''': 0, '''a.10''': 1, '''a.2''': 2}
         UpperCAmelCase_     :				List[str]  =     extract_submodules_state_dict(_SCREAMING_SNAKE_CASE  ,['''a.1''', '''a.2''']    )
         self.assertDictEqual(_SCREAMING_SNAKE_CASE  ,{'''a.1''': 0, '''a.2''': 2}    )

         UpperCAmelCase_     :				Dict  =     {'''a.1.a''': 0, '''a.10.a''': 1, '''a.2.a''': 2}
         UpperCAmelCase_     :				Optional[Any]  =     extract_submodules_state_dict(_SCREAMING_SNAKE_CASE  ,['''a.1''', '''a.2''']    )
         self.assertDictEqual(_SCREAMING_SNAKE_CASE  ,{'''a.1.a''': 0, '''a.2.a''': 2}    )





import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
     import torch

     from transformers import AutoModelForImageClassification

if is_vision_available():
     from transformers import AutoImageProcessor

@require_torch
@require_vision
class     __a(	unittest.TestCase  ):






    """simple docstring"""





    @slow
    def 				a__ ( self    )					->						List[str]:
         UpperCAmelCase_     :				Tuple  =     AutoImageProcessor.from_pretrained('''microsoft/dit-base-finetuned-rvlcdip'''    )
         UpperCAmelCase_     :				Union[str, Any]  =     AutoModelForImageClassification.from_pretrained('''microsoft/dit-base-finetuned-rvlcdip'''    )
         model.to(_SCREAMING_SNAKE_CASE    )

         from datasets import load_dataset

         UpperCAmelCase_     :				Optional[int]  =     load_dataset('''nielsr/rvlcdip-demo'''    )

         UpperCAmelCase_     :				Optional[Any]  =     dataset['''train'''][0]['''image'''].convert('''RGB'''    )

         UpperCAmelCase_     :				str  =     image_processor(_SCREAMING_SNAKE_CASE  ,return_tensors='''pt'''    ).to(_SCREAMING_SNAKE_CASE    )

         # forward pass
         with torch.no_grad():
              UpperCAmelCase_     :				Optional[int]  =     model(**_SCREAMING_SNAKE_CASE    )
              UpperCAmelCase_     :				List[Any]  =     outputs.logits

         UpperCAmelCase_     :				Tuple  =     torch.Size((1, 16)    )
         self.assertEqual(logits.shape  ,_SCREAMING_SNAKE_CASE    )

         UpperCAmelCase_     :				int  =     torch.tensor(
             [-0.41_58, -0.40_92, -0.43_47]  ,device=_SCREAMING_SNAKE_CASE  ,dtype=torch.float  ,)
         self.assertTrue(torch.allclose(logits[0, :3]  ,_SCREAMING_SNAKE_CASE  ,atol=1e-4    )    )