'''simple docstring'''






from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


_UpperCAmelCase						:    Any     					=					{
    '''configuration_ernie_m''': ['''ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ErnieMConfig'''],
}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase						:    Optional[Any]     					=					['''ErnieMTokenizer''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase						:    str     					=					[
       '''ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''ErnieMForMultipleChoice''',
       '''ErnieMForQuestionAnswering''',
       '''ErnieMForSequenceClassification''',
       '''ErnieMForTokenClassification''',
       '''ErnieMModel''',
       '''ErnieMPreTrainedModel''',
       '''ErnieMForInformationExtraction''',
   ]


if TYPE_CHECKING:
   from .configuration_ernie_m import ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieMConfig

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_ernie_m import ErnieMTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_ernie_m import (
          ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST,
          ErnieMForInformationExtraction,
          ErnieMForMultipleChoice,
          ErnieMForQuestionAnswering,
          ErnieMForSequenceClassification,
          ErnieMForTokenClassification,
          ErnieMModel,
          ErnieMPreTrainedModel,
      )


else:
   import sys

   _UpperCAmelCase						:    List[str]     					=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




'''simple docstring'''






import torch

from accelerate import Accelerator, DistributedType







class lowercase_					:







   """simple docstring"""

   def __enter__(			self			: Optional[int]    )     -> Union[str, Any]:
        if self.enabled:
             _A								=     self.model.no_sync()
             self.model_sync_obj.__enter__()

        return self
   def __exit__(			self			: Dict,   UpperCamelCase__			: Any,   UpperCamelCase__			: int,   UpperCamelCase__			: Any    )     -> str:
        if self.enabled:
             # Average all models on exit
             self._sync_and_avg_model_params()
             self.model_sync_obj.__exit__(UpperCamelCase__,   UpperCamelCase__,   UpperCamelCase__    )
   def __init__(			self			: Optional[int],   UpperCamelCase__			: Accelerator,   UpperCamelCase__			: torch.nn.Module,   UpperCamelCase__			: int,   UpperCamelCase__			: bool = True    )     -> str:
        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)'    )
        _A								=     enabled and accelerator.distributed_type != DistributedType.NO
        _A								=     0
        if self.enabled:
             _A								=     accelerator
             _A								=     model
             _A								=     local_sgd_steps
   def   __UpperCAmelCase      (			self			: Any    )     -> int:
        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   __UpperCAmelCase      (			self			: Dict    )     -> int:
        self.accelerator.wait_for_everyone()
        with self.accelerator.autocast():
             for param in self.model.parameters():
                  _A								=     self.accelerator.reduce(param.data,   reduction='mean'    )



