



"""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					(       snake_case_  :  Tuple   ,	snake_case_  :  Tuple   ,	snake_case_  :  str   ,	snake_case_  :  List[Any]=None   )							->	Optional[Any]:






   """simple docstring"""


   _lowerCAmelCase					     =		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.
      _lowerCAmelCase					     =		weight.view(torch.intaa   )
      _lowerCAmelCase					     =		"""bfloat16"""
   _lowerCAmelCase					     =		weight.cpu().numpy()
   _lowerCAmelCase					     =		os.path.join(snake_case_   ,	F"""{weight_name}.dat"""   )
   if index is not None:
      if dtype is None:
         _lowerCAmelCase					     =		str(array.dtype   )
      _lowerCAmelCase					     =		{"""dtype""": dtype, """shape""": list(array.shape   )}
   if array.ndim == 0:
      _lowerCAmelCase					     =		array[None]
   _lowerCAmelCase					     =		np.memmap(snake_case_   ,	dtype=array.dtype   ,	mode="""w+"""   ,	shape=array.shape   )
   _lowerCAmelCase					     =		array[:]
   file_array.flush()
   return index






def 	__UpperCAmelCase					(       snake_case_  :  Tuple   ,	snake_case_  :  Tuple   )							->	Any:






   """simple docstring"""


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

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

   _lowerCAmelCase					     =		np.memmap(snake_case_   ,	dtype=snake_case_   ,	shape=snake_case_   ,	mode="""r"""   )

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

   return weight






def 	__UpperCAmelCase					(       snake_case_  :  Any   ,	snake_case_  :  str   )							->	str:






   """simple docstring"""


   if index is None or len(snake_case_   ) == 0:
      # Nothing to save
      return

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

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






def 	__UpperCAmelCase					(       snake_case_  :  Union[str, os.PathLike]   ,	snake_case_  :  Dict[str, torch.Tensor]   )							->	Tuple:






   """simple docstring"""


   os.makedirs(snake_case_   ,	exist_ok=snake_case_   )
   _lowerCAmelCase					     =		{}
   for name, parameter in state_dict.items():
      _lowerCAmelCase					     =		offload_weight(snake_case_   ,	snake_case_   ,	snake_case_   ,	index=snake_case_   )

   # Update index
   save_offload_index(snake_case_   ,	snake_case_   )


class       __lowerCamelCase	(   __lowercase    ):







   def __init__(self       ,	lowerCamelCase       ,	lowerCamelCase				):




      '''simple docstring'''




      _lowerCAmelCase					     =		dataset
      _lowerCAmelCase					     =		prefix







   def __getitem__(self       ,	lowerCamelCase				):




      '''simple docstring'''




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







   def __iter__(self				):




      '''simple docstring'''




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







   def __len__(self				):




      '''simple docstring'''




      return len(self.dataset				)





class       __lowerCamelCase	(   __lowercase    ):







   def __init__(self       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase=None       ,	):




      '''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."""				)

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







   def __getitem__(self       ,	lowerCamelCase				):




      '''simple docstring'''




      if key in self.state_dict:
         return self.state_dict[key]
      _lowerCAmelCase					     =		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

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

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

      _lowerCAmelCase					     =		os.path.join(self.save_folder       ,	f"""{key}.dat"""				)
      return load_offloaded_weight(lowerCamelCase       ,	lowerCamelCase				)







   def __iter__(self				):




      '''simple docstring'''




      return iter(self.all_keys				)







   def __len__(self				):




      '''simple docstring'''




      return len(self.all_keys				)






def 	__UpperCAmelCase					(       snake_case_  :  Dict[str, torch.Tensor]   ,	snake_case_  :  List[str]   )							->	Optional[Any]:






   """simple docstring"""


   _lowerCAmelCase					     =		{}
   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




"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


SCREAMING_SNAKE_CASE				:     Union[str, Any]          =	{
    '''configuration_align''': [
        '''ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''AlignConfig''',
        '''AlignTextConfig''',
        '''AlignVisionConfig''',
    ],
    '''processing_align''': ['''AlignProcessor'''],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  SCREAMING_SNAKE_CASE				:     Any          =	[
      '''ALIGN_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''AlignModel''',
      '''AlignPreTrainedModel''',
      '''AlignTextModel''',
      '''AlignVisionModel''',
  ]

if TYPE_CHECKING:
  from .configuration_align import (
      ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP,
      AlignConfig,
      AlignTextConfig,
      AlignVisionConfig,
  )
  from .processing_align import AlignProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_align import (
        ALIGN_PRETRAINED_MODEL_ARCHIVE_LIST,
        AlignModel,
        AlignPreTrainedModel,
        AlignTextModel,
        AlignVisionModel,
    )

else:
  import sys

  SCREAMING_SNAKE_CASE				:     Any          =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)