"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


__SCREAMING_SNAKE_CASE								={
    "configuration_xlm": ["XLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLMConfig", "XLMOnnxConfig"],
    "tokenization_xlm": ["XLMTokenizer"],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __SCREAMING_SNAKE_CASE								=[
           "XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "XLMForMultipleChoice",
           "XLMForQuestionAnswering",
           "XLMForQuestionAnsweringSimple",
           "XLMForSequenceClassification",
           "XLMForTokenClassification",
           "XLMModel",
           "XLMPreTrainedModel",
           "XLMWithLMHeadModel",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __SCREAMING_SNAKE_CASE								=[
           "TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFXLMForMultipleChoice",
           "TFXLMForQuestionAnsweringSimple",
           "TFXLMForSequenceClassification",
           "TFXLMForTokenClassification",
           "TFXLMMainLayer",
           "TFXLMModel",
           "TFXLMPreTrainedModel",
           "TFXLMWithLMHeadModel",
       ]


if TYPE_CHECKING:
       from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
       from .tokenization_xlm import XLMTokenizer

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_xlm import (
                  XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  XLMForMultipleChoice,
                  XLMForQuestionAnswering,
                  XLMForQuestionAnsweringSimple,
                  XLMForSequenceClassification,
                  XLMForTokenClassification,
                  XLMModel,
                  XLMPreTrainedModel,
                  XLMWithLMHeadModel,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_xlm import (
                  TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFXLMForMultipleChoice,
                  TFXLMForQuestionAnsweringSimple,
                  TFXLMForSequenceClassification,
                  TFXLMForTokenClassification,
                  TFXLMMainLayer,
                  TFXLMModel,
                  TFXLMPreTrainedModel,
                  TFXLMWithLMHeadModel,
              )

else:
       import sys

       __SCREAMING_SNAKE_CASE								=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


"""simple docstring"""







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  UpperCamelCase	(    nn.Module						):
   def __init__(    self    )			->   Any:


        '''simple docstring'''


        super().__init__()
        lowercase_			:       Dict									=			nn.Linear(3  ,4    )
        lowercase_			:       Union[str, Any]									=			nn.BatchNormad(4    )
        lowercase_			:       Union[str, Any]									=			nn.Linear(4  ,5    )







   def     _UpperCAmelCase     (    self  ,__UpperCamelCase    )			->   int:


        '''simple docstring'''


        return self.lineara(self.batchnorm(self.lineara(__UpperCamelCase    )    )    )







class  UpperCamelCase	(    unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       int									=			ModelForTest()
        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,model.state_dict()    )
             lowercase_			:       str									=			os.path.join(__UpperCamelCase  ,'index.json'    )
             self.assertTrue(os.path.isfile(__UpperCamelCase    )    )
             # TODO: add tests on what is inside the index

             for key in ["linear1.weight", "linear1.bias", "linear2.weight", "linear2.bias"]:
                  lowercase_			:       int									=			os.path.join(__UpperCamelCase  ,f'''{key}.dat'''    )
                  self.assertTrue(os.path.isfile(__UpperCamelCase    )    )
                  # TODO: add tests on the fact weights are properly loaded
   def     _UpperCAmelCase     (    self    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       List[str]									=			[torch.floataa, torch.floataa, torch.bfloataa]

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

                  lowercase_			:       Optional[Any]									=			load_offloaded_weight(__UpperCamelCase  ,index['weight']    )
                  self.assertTrue(torch.equal(__UpperCamelCase  ,__UpperCamelCase    )    )
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        lowercase_			:       List[str]									=			ModelForTest()
        lowercase_			:       Optional[int]									=			model.state_dict()
        lowercase_			:       int									=			{k: v for k, v in state_dict.items() if 'linear2' not in k}
        lowercase_			:       Optional[Any]									=			{k: v for k, v in state_dict.items() if 'linear2' in k}

        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,__UpperCamelCase    )
             lowercase_			:       Tuple									=			OffloadedWeightsLoader(state_dict=__UpperCamelCase  ,save_folder=__UpperCamelCase    )

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

        lowercase_			:       int									=			{k: v for k, v in state_dict.items() if 'weight' in k}
        lowercase_			:       int									=			{k: v for k, v in state_dict.items() if 'weight' not in k}

        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,__UpperCamelCase    )
             lowercase_			:       int									=			OffloadedWeightsLoader(state_dict=__UpperCamelCase  ,save_folder=__UpperCamelCase    )

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

        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,__UpperCamelCase    )
             # Duplicates are removed
             lowercase_			:       Tuple									=			OffloadedWeightsLoader(state_dict=__UpperCamelCase  ,save_folder=__UpperCamelCase    )

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







   def     _UpperCAmelCase     (    self    )			->   Union[str, Any]:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			{'a.1': 0, 'a.10': 1, 'a.2': 2}
        lowercase_			:       int									=			extract_submodules_state_dict(__UpperCamelCase  ,['a.1', 'a.2']    )
        self.assertDictEqual(__UpperCamelCase  ,{'a.1': 0, 'a.2': 2}    )

        lowercase_			:       Union[str, Any]									=			{'a.1.a': 0, 'a.10.a': 1, 'a.2.a': 2}
        lowercase_			:       Optional[int]									=			extract_submodules_state_dict(__UpperCamelCase  ,['a.1', 'a.2']    )
        self.assertDictEqual(__UpperCamelCase  ,{'a.1.a': 0, 'a.2.a': 2}    )

