





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    ):


				def __init__(		self							):
								super().__init__()
								lowerCamelCase	 =nn.Linear(3   , 4							)
								lowerCamelCase	 =nn.BatchNormad(4							)
								lowerCamelCase	 =nn.Linear(4   , 5							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								return self.lineara(self.batchnorm(self.lineara(UpperCAmelCase_							)							)							)




class       __A     (						unittest.TestCase    ):


				def _snake_case       (		self							):
								lowerCamelCase	 =ModelForTest()
								with TemporaryDirectory() as tmp_dir:
												offload_state_dict(UpperCAmelCase_   , model.state_dict()							)
												lowerCamelCase	 =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"]:
																lowerCamelCase	 =os.path.join(UpperCAmelCase_   , f"""{key}.dat"""							)
																self.assertTrue(os.path.isfile(UpperCAmelCase_							)							)
																# TODO: add tests on the fact weights are properly loaded


				def _snake_case       (		self							):
								lowerCamelCase	 =[torch.floataa, torch.floataa, torch.bfloataa]

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

																lowerCamelCase	 =load_offloaded_weight(UpperCAmelCase_   , index["""weight"""]							)
																self.assertTrue(torch.equal(UpperCAmelCase_   , UpperCAmelCase_							)							)


				def _snake_case       (		self							):
								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(UpperCAmelCase_   , UpperCAmelCase_							)
												lowerCamelCase	 =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]							)							)

								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(UpperCAmelCase_   , UpperCAmelCase_							)
												lowerCamelCase	 =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
												lowerCamelCase	 =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 _snake_case       (		self							):
								lowerCamelCase	 ={"""a.1""": 0, """a.10""": 1, """a.2""": 2}
								lowerCamelCase	 =extract_submodules_state_dict(UpperCAmelCase_   , ["""a.1""", """a.2"""]							)
								self.assertDictEqual(UpperCAmelCase_   , {"""a.1""": 0, """a.2""": 2}							)

								lowerCamelCase	 ={"""a.1.a""": 0, """a.10.a""": 1, """a.2.a""": 2}
								lowerCamelCase	 =extract_submodules_state_dict(UpperCAmelCase_   , ["""a.1""", """a.2"""]							)
								self.assertDictEqual(UpperCAmelCase_   , {"""a.1.a""": 0, """a.2.a""": 2}							)












from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch

@require_torch
@require_tf
class       __A     (						a    ):


				def _snake_case       (		self							):
								lowerCamelCase	 =SMALL_MODEL_IDENTIFIER
								lowerCamelCase	 ="""pt"""
								lowerCamelCase	 ="""tf"""


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =AutoModel.from_pretrained(self.test_model							)
								model_pt.save_pretrained(UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =TFAutoModel.from_pretrained(self.test_model   , from_pt=UpperCAmelCase_							)
								model_tf.save_pretrained(UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 ="""mock_framework"""

								# Framework provided - return whatever the user provides
								lowerCamelCase	 =FeaturesManager.determine_framework(self.test_model   , UpperCAmelCase_							)
								self.assertEqual(UpperCAmelCase_   , UpperCAmelCase_							)

								# Local checkpoint and framework provided - return provided framework
								# PyTorch checkpoint
								with TemporaryDirectory() as local_pt_ckpt:
												self._setup_pt_ckpt(UpperCAmelCase_							)
												lowerCamelCase	 =FeaturesManager.determine_framework(UpperCAmelCase_   , UpperCAmelCase_							)
												self.assertEqual(UpperCAmelCase_   , UpperCAmelCase_							)

								# TensorFlow checkpoint
								with TemporaryDirectory() as local_tf_ckpt:
												self._setup_tf_ckpt(UpperCAmelCase_							)
												lowerCamelCase	 =FeaturesManager.determine_framework(UpperCAmelCase_   , UpperCAmelCase_							)
												self.assertEqual(UpperCAmelCase_   , UpperCAmelCase_							)


				def _snake_case       (		self							):
								# PyTorch checkpoint
								with TemporaryDirectory() as local_pt_ckpt:
												self._setup_pt_ckpt(UpperCAmelCase_							)
												lowerCamelCase	 =FeaturesManager.determine_framework(UpperCAmelCase_							)
												self.assertEqual(UpperCAmelCase_   , self.framework_pt							)

								# TensorFlow checkpoint
								with TemporaryDirectory() as local_tf_ckpt:
												self._setup_tf_ckpt(UpperCAmelCase_							)
												lowerCamelCase	 =FeaturesManager.determine_framework(UpperCAmelCase_							)
												self.assertEqual(UpperCAmelCase_   , self.framework_tf							)

								# Invalid local checkpoint
								with TemporaryDirectory() as local_invalid_ckpt:
												with self.assertRaises(UpperCAmelCase_							):
																lowerCamelCase	 =FeaturesManager.determine_framework(UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =MagicMock(return_value=UpperCAmelCase_							)
								with patch("""transformers.onnx.features.is_tf_available"""   , UpperCAmelCase_							):
												lowerCamelCase	 =FeaturesManager.determine_framework(self.test_model							)
												self.assertEqual(UpperCAmelCase_   , self.framework_pt							)

								# PyTorch not in environment -> use TensorFlow
								lowerCamelCase	 =MagicMock(return_value=UpperCAmelCase_							)
								with patch("""transformers.onnx.features.is_torch_available"""   , UpperCAmelCase_							):
												lowerCamelCase	 =FeaturesManager.determine_framework(self.test_model							)
												self.assertEqual(UpperCAmelCase_   , self.framework_tf							)

								# Both in environment -> use PyTorch
								lowerCamelCase	 =MagicMock(return_value=UpperCAmelCase_							)
								lowerCamelCase	 =MagicMock(return_value=UpperCAmelCase_							)
								with patch("""transformers.onnx.features.is_tf_available"""   , UpperCAmelCase_							), patch(
								    """transformers.onnx.features.is_torch_available"""   , UpperCAmelCase_							):
												lowerCamelCase	 =FeaturesManager.determine_framework(self.test_model							)
												self.assertEqual(UpperCAmelCase_   , self.framework_pt							)

								# Both not in environment -> raise error
								lowerCamelCase	 =MagicMock(return_value=UpperCAmelCase_							)
								lowerCamelCase	 =MagicMock(return_value=UpperCAmelCase_							)
								with patch("""transformers.onnx.features.is_tf_available"""   , UpperCAmelCase_							), patch(
								    """transformers.onnx.features.is_torch_available"""   , UpperCAmelCase_							):
												with self.assertRaises(UpperCAmelCase_							):
																lowerCamelCase	 =FeaturesManager.determine_framework(self.test_model							)





