




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 :  int  )       ->     Optional[Any]:

       """simple docstring"""






       super().__init__()
       lowerCamelCase_							=   nn.Linear(3						,							4  )
       lowerCamelCase_							=   nn.BatchNormad(4  )
       lowerCamelCase_							=   nn.Linear(4						,							5  )





 def     a__			(    self :  int						,							A_ :  Optional[int]  )       ->     List[Any]:

       """simple docstring"""






       return self.lineara(self.batchnorm(self.lineara(_snake_case  )  )  )







class   A(						unittest.TestCase       ):



 '''simple docstring'''


 def     a__			(    self :  str  )       ->     List[str]:

       """simple docstring"""






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

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

 def     a__			(    self :  List[str]  )       ->     int:

       """simple docstring"""






       lowerCamelCase_							=   [torch.floataa, torch.floataa, torch.bfloataa]

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

                   lowerCamelCase_							=   load_offloaded_weight(_snake_case						,							index['weight']  )
                   self.assertTrue(torch.equal(_snake_case						,							_snake_case  )  )

 def     a__			(    self :  Optional[Any]  )       ->     Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=   ModelForTest()
       lowerCamelCase_							=   model.state_dict()
       lowerCamelCase_							=   {k: v for k, v in state_dict.items() if "linear2" not in k}
       lowerCamelCase_							=   {k: v for k, v in state_dict.items() if "linear2" in k}

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

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

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

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

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

       with TemporaryDirectory() as tmp_dir:
             offload_state_dict(_snake_case						,							_snake_case  )
             # Duplicates are removed
             lowerCamelCase_							=   OffloadedWeightsLoader(state_dict=_snake_case						,							save_folder=_snake_case  )

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





 def     a__			(    self :  List[str]  )       ->     Any:

       """simple docstring"""






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

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









import os
import time

import numpy as np
import onnxruntime as ort


lowerCamelCase    :		int			   =       "1"
lowerCamelCase    :		int			   =       "0"
lowerCamelCase    :		Union[str, Any]			   =       "1"

lowerCamelCase    :		List[Any]			   =       ort.SessionOptions()
lowerCamelCase    :		Optional[Any]			   =       ort.GraphOptimizationLevel.ORT_DISABLE_ALL
print("Create inference session...")
lowerCamelCase    :		Union[str, Any]			   =       ["TensorrtExecutionProvider", "CUDAExecutionProvider"]
lowerCamelCase    :		Tuple			   =       ort.InferenceSession("model.onnx", sess_options=sess_opt, providers=execution_provider)
lowerCamelCase    :		List[Any]			   =       ort.RunOptions()

lowerCamelCase    :		List[str]			   =       128
lowerCamelCase    :		List[Any]			   =       1
lowerCamelCase    :		Union[str, Any]			   =       np.ones((batch, sequence), dtype=np.intaa)
lowerCamelCase    :		Dict			   =       np.ones((batch, sequence), dtype=np.intaa)
lowerCamelCase    :		Optional[Any]			   =       np.ones((batch, sequence), dtype=np.intaa)

print("Warm up phase...")
sess.run(
    None,
    {
        sess.get_inputs()[0].name: input_ids,
        sess.get_inputs()[1].name: attention_mask,
        sess.get_inputs()[2].name: token_type_ids,
    },
    run_options=run_opt,
)

print("Start inference...")
lowerCamelCase    :		int			   =       time.time()
lowerCamelCase    :		Dict			   =       2_000
lowerCamelCase    :		Any			   =       {}
for iter in range(max_iters):
 lowerCamelCase    :		Union[str, Any]			   =       sess.run(
     None,
     {
         sess.get_inputs()[0].name: input_ids,
         sess.get_inputs()[1].name: attention_mask,
         sess.get_inputs()[2].name: token_type_ids,
     },
     run_options=run_opt,
 )
print("Average Inference Time = {:.3f} ms".format((time.time() - start_time) * 1_000 / max_iters))



