





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        _lowerCAmelCase    (			nn.Module							):


  '''simple docstring'''







  def __init__(      self      :					str							):


    '''simple docstring'''





    super().__init__()
    _snake_case   :				Dict														=	nn.Linear(3      ,				4							)
    _snake_case   :				int														=	nn.BatchNormad(4							)
    _snake_case   :				Tuple														=	nn.Linear(4      ,				5							)







  def UpperCamelCase_ (      self      :					Dict      ,				UpperCamelCase      :					int							):


    '''simple docstring'''





    return self.lineara(self.batchnorm(self.lineara(a_							)							)							)




class        _lowerCAmelCase    (			unittest.TestCase							):


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





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

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







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				List[Any]														=	[torch.floataa, torch.floataa, torch.bfloataa]

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

        _snake_case   :				List[str]														=	load_offloaded_weight(a_      ,				index['weight']							)
        self.assertTrue(torch.equal(a_      ,				a_							)							)







  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





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

    with TemporaryDirectory() as tmp_dir:
      offload_state_dict(a_      ,				a_							)
      _snake_case   :				Tuple														=	OffloadedWeightsLoader(state_dict=a_      ,				save_folder=a_							)

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

    _snake_case   :				Dict														=	{k: v for k, v in state_dict.items() if """weight""" in k}
    _snake_case   :				str														=	{k: v for k, v in state_dict.items() if """weight""" not in k}

    with TemporaryDirectory() as tmp_dir:
      offload_state_dict(a_      ,				a_							)
      _snake_case   :				Optional[int]														=	OffloadedWeightsLoader(state_dict=a_      ,				save_folder=a_							)

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

    with TemporaryDirectory() as tmp_dir:
      offload_state_dict(a_      ,				a_							)
      # Duplicates are removed
      _snake_case   :				Union[str, Any]														=	OffloadedWeightsLoader(state_dict=a_      ,				save_folder=a_							)

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







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	{"""a.1""": 0, """a.10""": 1, """a.2""": 2}
    _snake_case   :				Optional[Any]														=	extract_submodules_state_dict(a_      ,				['a.1', 'a.2']							)
    self.assertDictEqual(a_      ,				{'a.1': 0, 'a.2': 2}							)

    _snake_case   :				List[Any]														=	{"""a.1.a""": 0, """a.10.a""": 1, """a.2.a""": 2}
    _snake_case   :				Tuple														=	extract_submodules_state_dict(a_      ,				['a.1', 'a.2']							)
    self.assertDictEqual(a_      ,				{'a.1.a': 0, 'a.2.a': 2}							)











def     lowerCamelCase_	(      lowerCAmelCase:	bytes				)->		str:
  return "".join([hex(lowerCAmelCase				)[2:].zfill(2				).upper() for byte in list(lowerCAmelCase				)]				)


def     lowerCamelCase_	(      lowerCAmelCase:	str				)->		bytes:
  # Check data validity, following RFC3548
  # https://www.ietf.org/rfc/rfc3548.txt
  if (len(lowerCAmelCase				) % 2) != 0:
    raise ValueError(
        'Base16 encoded data is invalid:\nData does not have an even number of hex digits.'				)
  # Check the character set - the standard base16 alphabet
  # is uppercase according to RFC3548 section 6
  if not set(lowerCAmelCase				) <= set('0123456789ABCDEF'				):
    raise ValueError(
        'Base16 encoded data is invalid:\nData is not uppercase hex or it contains invalid characters.'				)
  # For every two hexadecimal digits (= a byte), turn it into an integer.
  # Then, string the result together into bytes, and return it.
  return bytes(int(data[i] + data[i + 1]    ,							16				) for i in range(0    ,							len(lowerCAmelCase				)    ,							2				)				)


if __name__ == "__main__":
      import doctest

      doctest.testmod()




