
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowerCamelCase__        =						{
    """configuration_esm""": ["""ESM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """EsmConfig"""],
    """tokenization_esm""": ["""EsmTokenizer"""],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase__        =						[
					    """ESM_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """EsmForMaskedLM""",
					    """EsmForSequenceClassification""",
					    """EsmForTokenClassification""",
					    """EsmModel""",
					    """EsmPreTrainedModel""",
					]
					lowerCamelCase__        =						["""EsmForProteinFolding""", """EsmFoldPreTrainedModel"""]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase__        =						[
					    """TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """TFEsmForMaskedLM""",
					    """TFEsmForSequenceClassification""",
					    """TFEsmForTokenClassification""",
					    """TFEsmModel""",
					    """TFEsmPreTrainedModel""",
					]

if TYPE_CHECKING:
					from .configuration_esm import ESM_PRETRAINED_CONFIG_ARCHIVE_MAP, EsmConfig
					from .tokenization_esm import EsmTokenizer

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_esm import (
										    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
										    EsmForMaskedLM,
										    EsmForSequenceClassification,
										    EsmForTokenClassification,
										    EsmModel,
										    EsmPreTrainedModel,
										)
										from .modeling_esmfold import EsmFoldPreTrainedModel, EsmForProteinFolding

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_esm import (
										    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFEsmForMaskedLM,
										    TFEsmForSequenceClassification,
										    TFEsmForTokenClassification,
										    TFEsmModel,
										    TFEsmPreTrainedModel,
										)


else:
					import sys

					lowerCamelCase__        =						_LazyModule(__name__, globals()["""__file__"""], _import_structure)

import torch

from accelerate import Accelerator, DistributedType







class 	A__				:






				def __enter__(		self   : Optional[Any]  ):






							'''simple docstring'''







							if self.enabled:
										lowerCAmelCase__ : Any   =						self.model.no_sync()
										self.model_sync_obj.__enter__()

							return self






				def __exit__(		self   : int       ,	a   : int       ,	a   : Optional[Any]       ,	a   : str  ):






							'''simple docstring'''







							if self.enabled:
										# Average all models on exit
										self._sync_and_avg_model_params()
										self.model_sync_obj.__exit__(a       ,	a       ,	a  )






				def __init__(		self   : List[str]       ,	a   : Accelerator       ,	a   : torch.nn.Module       ,	a   : int       ,	a   : bool = True  ):






							'''simple docstring'''







							if accelerator.distributed_type not in [
							    DistributedType.NO,
							    DistributedType.MULTI_CPU,
							    DistributedType.MULTI_GPU,
							]:
										raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'  )
							lowerCAmelCase__ : Union[str, Any]   =						enabled and accelerator.distributed_type != DistributedType.NO
							lowerCAmelCase__ : str   =						0
							if self.enabled:
										lowerCAmelCase__ : Dict   =						accelerator
										lowerCAmelCase__ : Dict   =						model
										lowerCAmelCase__ : Optional[int]   =						local_sgd_steps






				def 				_lowerCamelCase						(		self   : Optional[Any]  ):






							'''simple docstring'''







							self.num_steps += 1
							if not self.enabled:
										return

							if self.num_steps % self.local_sgd_steps == 0:
										self._sync_and_avg_model_params()






				def 				_lowerCamelCase						(		self   : int  ):






							'''simple docstring'''







							self.accelerator.wait_for_everyone()
							with self.accelerator.autocast():
										for param in self.model.parameters():
													lowerCAmelCase__ : Dict   =						self.accelerator.reduce(param.data       ,	reduction='mean'  )