



"""simple docstring"""


import safs







class  A_     (	safs.SaFileSystem							):



						"""simple docstring"""







"""simple docstring"""




import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_=None     ):
					SCREAMING_SNAKE_CASE           =      None
					# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
					if str(weight.dtype     ) == "torch.bfloat16":
										# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
										SCREAMING_SNAKE_CASE           =      weight.view(torch.intaa     )
										SCREAMING_SNAKE_CASE           =      'bfloat16'
					SCREAMING_SNAKE_CASE           =      weight.cpu().numpy()
					SCREAMING_SNAKE_CASE           =      os.path.join(SCREAMING_SNAKE_CASE_,	f'''{weight_name}.dat'''     )
					if index is not None:
										if dtype is None:
															SCREAMING_SNAKE_CASE           =      str(array.dtype     )
										SCREAMING_SNAKE_CASE           =      {'dtype': dtype, 'shape': list(array.shape     )}
					if array.ndim == 0:
										SCREAMING_SNAKE_CASE           =      array[None]
					SCREAMING_SNAKE_CASE           =      np.memmap(SCREAMING_SNAKE_CASE_,	dtype=array.dtype,	mode='w+',	shape=array.shape     )
					SCREAMING_SNAKE_CASE           =      array[:]
					file_array.flush()
					return index

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					SCREAMING_SNAKE_CASE           =      tuple(weight_info['shape']     )
					if shape == ():
										# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
										SCREAMING_SNAKE_CASE           =      (1,)

					SCREAMING_SNAKE_CASE           =      weight_info['dtype']
					if dtype == "bfloat16":
										# NumPy does not support bfloat16 so this was saved as a int16
										SCREAMING_SNAKE_CASE           =      'int16'

					SCREAMING_SNAKE_CASE           =      np.memmap(SCREAMING_SNAKE_CASE_,	dtype=SCREAMING_SNAKE_CASE_,	shape=SCREAMING_SNAKE_CASE_,	mode='r'     )

					if len(weight_info['shape']     ) == 0:
										SCREAMING_SNAKE_CASE           =      weight[0]
					SCREAMING_SNAKE_CASE           =      torch.tensor(SCREAMING_SNAKE_CASE_     )
					if weight_info["dtype"] == "bfloat16":
										SCREAMING_SNAKE_CASE           =      weight.view(torch.bfloataa     )

					return weight

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					if index is None or len(SCREAMING_SNAKE_CASE_     ) == 0:
										# Nothing to save
										return

					SCREAMING_SNAKE_CASE           =      os.path.join(SCREAMING_SNAKE_CASE_,	'index.json'     )
					if os.path.isfile(SCREAMING_SNAKE_CASE_     ):
										with open(SCREAMING_SNAKE_CASE_,	'r',	encoding='utf-8'     ) as f:
															SCREAMING_SNAKE_CASE           =      json.load(SCREAMING_SNAKE_CASE_     )
					else:
										SCREAMING_SNAKE_CASE           =      {}
					current_index.update(SCREAMING_SNAKE_CASE_     )

					with open(SCREAMING_SNAKE_CASE_,	'w',	encoding='utf-8'     ) as f:
										json.dump(SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	indent=2     )

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					os.makedirs(SCREAMING_SNAKE_CASE_,	exist_ok=SCREAMING_SNAKE_CASE_     )
					SCREAMING_SNAKE_CASE           =      {}
					for name, parameter in state_dict.items():
										SCREAMING_SNAKE_CASE           =      offload_weight(SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	index=SCREAMING_SNAKE_CASE_     )

					# Update index
					save_offload_index(SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     )




class 	UpperCamelCase      (			__magic_name__ ):

				"""simple docstring"""



				def __init__(							self						,      lowercase__						,      lowercase__     )       ->       Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      dataset
									SCREAMING_SNAKE_CASE           =      prefix


				def __getitem__(							self						,      lowercase__     )       ->       Dict:
									"""simple docstring"""







									return self.dataset[f'''{self.prefix}{key}''']


				def __iter__(							self     )       ->       List[str]:
									"""simple docstring"""







									return iter([key for key in self.dataset if key.startswith(self.prefix     )]     )




				def __len__(							self     )       ->       Optional[Any]:
									"""simple docstring"""







									return len(self.dataset     )






class 	UpperCamelCase      (			__magic_name__ ):

				"""simple docstring"""



				def __init__(							self						,      lowercase__ = None						,      lowercase__ = None						,      lowercase__ = None						,      lowercase__=None						,      )       ->       Optional[Any]:
									"""simple docstring"""







									if state_dict is None and save_folder is None:
														raise ValueError('Need either a `state_dict` or a `save_folder` containing offloaded weights.'     )

									SCREAMING_SNAKE_CASE           =      {} if state_dict is None else state_dict
									SCREAMING_SNAKE_CASE           =      save_folder
									if index is None and save_folder is not None:
														with open(os.path.join(lowercase__						,      'index.json'     )     ) as f:
																			SCREAMING_SNAKE_CASE           =      json.load(lowercase__     )
									SCREAMING_SNAKE_CASE           =      {} if index is None else index
									SCREAMING_SNAKE_CASE           =      list(self.state_dict.keys()     )
									self.all_keys.extend([key for key in self.index if key not in self.all_keys]     )
									SCREAMING_SNAKE_CASE           =      device


				def __getitem__(							self						,      lowercase__     )       ->       Union[str, Any]:
									"""simple docstring"""







									if key in self.state_dict:
														return self.state_dict[key]
									SCREAMING_SNAKE_CASE           =      self.index[key]
									if weight_info.get('safetensors_file'     ) is not None:
														if not is_safetensors_available():
																			raise ImportError('These offloaded weights require the use of safetensors: `pip install safetensors`.'     )

														from safetensors import safe_open

														SCREAMING_SNAKE_CASE           =      'cpu' if self.device is None else self.device
														with safe_open(weight_info['safetensors_file']						,      framework='pt'						,      device=lowercase__     ) as f:
																			SCREAMING_SNAKE_CASE           =      f.get_tensor(weight_info.get('weight_name'						,      lowercase__     )     )

														if "dtype" in weight_info:
																			return tensor.to(getattr(lowercase__						,      weight_info['dtype']     )     )
														else:
																			return tensor

									SCREAMING_SNAKE_CASE           =      os.path.join(self.save_folder						,      f'''{key}.dat'''     )
									return load_offloaded_weight(lowercase__						,      lowercase__     )


				def __iter__(							self     )       ->       Any:
									"""simple docstring"""







									return iter(self.all_keys     )




				def __len__(							self     )       ->       Optional[Any]:
									"""simple docstring"""







									return len(self.all_keys     )






def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					SCREAMING_SNAKE_CASE           =      {}
					for module_name in submodule_names:
										# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
										# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
										result.update(
										    {
										        key: param
										        for key, param in state_dict.items()
										        if key == module_name or key.startswith(module_name + '.'     )
										    }     )
					return result
