
'''simple docstring'''







import contextlib
import copy
import random
from typing import Any, Dict, Iterable, Optional, Union

import numpy as np
import torch

from .utils import deprecate, is_transformers_available


if is_transformers_available():
					import transformers







def 				SCREAMING_SNAKE_CASE__	(   __A	)       ->	Optional[int]:
	random.seed(__A	)
	np.random.seed(__A	)
	torch.manual_seed(__A	)
	torch.cuda.manual_seed_all(__A	)
	# ^^ safe to call this function even if cuda is not available





class 						__UpperCAmelCase      :


							def __init__(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_ = 0.9999    ,   lowerCAmelCase_ = 0.0    ,   lowerCAmelCase_ = 0    ,   lowerCAmelCase_ = False    ,   lowerCAmelCase_ = 1.0    ,   lowerCAmelCase_ = 2 / 3    ,   lowerCAmelCase_ = None    ,   lowerCAmelCase_ = None    ,   **lowerCAmelCase_    ,   ):





								"""simple docstring"""






								if isinstance(lowerCAmelCase_    ,   torch.nn.Module  ):
									_snake_case   						=		(
									    'Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. '
									    'Please pass the parameters of the module instead.'
									)
									deprecate(
									    'passing a `torch.nn.Module` to `ExponentialMovingAverage`'    ,   '1.0.0'    ,   lowerCAmelCase_    ,   standard_warn=lowerCAmelCase_    ,   )
									_snake_case   						=		parameters.parameters()

									# set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility
									_snake_case   						=		True

								if kwargs.get('max_value'    ,   lowerCAmelCase_  ) is not None:
									_snake_case   						=		'The `max_value` argument is deprecated. Please use `decay` instead.'
									deprecate('max_value'    ,   '1.0.0'    ,   lowerCAmelCase_    ,   standard_warn=lowerCAmelCase_  )
									_snake_case   						=		kwargs['max_value']

								if kwargs.get('min_value'    ,   lowerCAmelCase_  ) is not None:
									_snake_case   						=		'The `min_value` argument is deprecated. Please use `min_decay` instead.'
									deprecate('min_value'    ,   '1.0.0'    ,   lowerCAmelCase_    ,   standard_warn=lowerCAmelCase_  )
									_snake_case   						=		kwargs['min_value']

								_snake_case   						=		list(lowerCAmelCase_  )
								_snake_case   						=		[p.clone().detach() for p in parameters]

								if kwargs.get('device'    ,   lowerCAmelCase_  ) is not None:
									_snake_case   						=		'The `device` argument is deprecated. Please use `to` instead.'
									deprecate('device'    ,   '1.0.0'    ,   lowerCAmelCase_    ,   standard_warn=lowerCAmelCase_  )
									self.to(device=kwargs['device']  )

								_snake_case   						=		None

								_snake_case   						=		decay
								_snake_case   						=		min_decay
								_snake_case   						=		update_after_step
								_snake_case   						=		use_ema_warmup
								_snake_case   						=		inv_gamma
								_snake_case   						=		power
								_snake_case   						=		0
								_snake_case   						=		None  # set in `step()`

								_snake_case   						=		model_cls
								_snake_case   						=		model_config


							@classmethod
							def  lowerCamelCase				(       cls    ,   lowerCAmelCase_    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case						, _snake_case   						=		model_cls.load_config(lowerCAmelCase_    ,   return_unused_kwargs=lowerCAmelCase_  )
								_snake_case   						=		model_cls.from_pretrained(lowerCAmelCase_  )

								_snake_case   						=		cls(model.parameters()    ,   model_cls=lowerCAmelCase_    ,   model_config=model.config  )

								ema_model.load_state_dict(lowerCAmelCase_  )
								return ema_model


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if self.model_cls is None:
									raise ValueError('`save_pretrained` can only be used if `model_cls` was defined at __init__.'  )

								if self.model_config is None:
									raise ValueError('`save_pretrained` can only be used if `model_config` was defined at __init__.'  )

								_snake_case   						=		self.model_cls.from_config(self.model_config  )
								_snake_case   						=		self.state_dict()
								state_dict.pop('shadow_params'    ,   lowerCAmelCase_  )

								model.register_to_config(**lowerCAmelCase_  )
								self.copy_to(model.parameters()  )
								model.save_pretrained(lowerCAmelCase_  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		max(0    ,   optimization_step - self.update_after_step - 1  )

								if step <= 0:
									return 0.0

								if self.use_ema_warmup:
									_snake_case   						=		1 - (1 + step / self.inv_gamma) ** -self.power
								else:
									_snake_case   						=		(1 + step) / (10 + step)

								_snake_case   						=		min(lowerCAmelCase_    ,   self.decay  )
								# make sure decay is not smaller than min_decay
								_snake_case   						=		max(lowerCAmelCase_    ,   self.min_decay  )
								return cur_decay_value


							@torch.no_grad()
							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if isinstance(lowerCAmelCase_    ,   torch.nn.Module  ):
									_snake_case   						=		(
									    'Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. '
									    'Please pass the parameters of the module instead.'
									)
									deprecate(
									    'passing a `torch.nn.Module` to `ExponentialMovingAverage.step`'    ,   '1.0.0'    ,   lowerCAmelCase_    ,   standard_warn=lowerCAmelCase_    ,   )
									_snake_case   						=		parameters.parameters()

								_snake_case   						=		list(lowerCAmelCase_  )

								self.optimization_step += 1

								# Compute the decay factor for the exponential moving average.
								_snake_case   						=		self.get_decay(self.optimization_step  )
								_snake_case   						=		decay
								_snake_case   						=		1 - decay

								_snake_case   						=		contextlib.nullcontext
								if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
									import deepspeed

								for s_param, param in zip(self.shadow_params    ,   lowerCAmelCase_  ):
									if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
										_snake_case   						=		deepspeed.zero.GatheredParameters(lowerCAmelCase_    ,   modifier_rank=lowerCAmelCase_  )

									with context_manager():
										if param.requires_grad:
											s_param.sub_(one_minus_decay * (s_param - param)  )
										else:
											s_param.copy_(lowerCAmelCase_  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		list(lowerCAmelCase_  )
								for s_param, param in zip(self.shadow_params    ,   lowerCAmelCase_  ):
									param.data.copy_(s_param.to(param.device  ).data  )


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





								"""simple docstring"""






								_snake_case   						=		[
								    p.to(device=lowerCAmelCase_    ,   dtype=lowerCAmelCase_  ) if p.is_floating_point() else p.to(device=lowerCAmelCase_  )
								    for p in self.shadow_params
								]


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								return {
								    "decay": self.decay,
								    "min_decay": self.min_decay,
								    "optimization_step": self.optimization_step,
								    "update_after_step": self.update_after_step,
								    "use_ema_warmup": self.use_ema_warmup,
								    "inv_gamma": self.inv_gamma,
								    "power": self.power,
								    "shadow_params": self.shadow_params,
								}


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		[param.detach().cpu().clone() for param in parameters]


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if self.temp_stored_params is None:
									raise RuntimeError('This ExponentialMovingAverage has no `store()`ed weights ' 'to `restore()`'  )
								for c_param, param in zip(self.temp_stored_params    ,   lowerCAmelCase_  ):
									param.data.copy_(c_param.data  )

								# Better memory-wise.
								_snake_case   						=		None




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		copy.deepcopy(lowerCAmelCase_  )

								_snake_case   						=		state_dict.get('decay'    ,   self.decay  )
								if self.decay < 0.0 or self.decay > 1.0:
									raise ValueError('Decay must be between 0 and 1'  )

								_snake_case   						=		state_dict.get('min_decay'    ,   self.min_decay  )
								if not isinstance(self.min_decay    ,   lowerCAmelCase_  ):
									raise ValueError('Invalid min_decay'  )

								_snake_case   						=		state_dict.get('optimization_step'    ,   self.optimization_step  )
								if not isinstance(self.optimization_step    ,   lowerCAmelCase_  ):
									raise ValueError('Invalid optimization_step'  )

								_snake_case   						=		state_dict.get('update_after_step'    ,   self.update_after_step  )
								if not isinstance(self.update_after_step    ,   lowerCAmelCase_  ):
									raise ValueError('Invalid update_after_step'  )

								_snake_case   						=		state_dict.get('use_ema_warmup'    ,   self.use_ema_warmup  )
								if not isinstance(self.use_ema_warmup    ,   lowerCAmelCase_  ):
									raise ValueError('Invalid use_ema_warmup'  )

								_snake_case   						=		state_dict.get('inv_gamma'    ,   self.inv_gamma  )
								if not isinstance(self.inv_gamma    ,   (float, int)  ):
									raise ValueError('Invalid inv_gamma'  )

								_snake_case   						=		state_dict.get('power'    ,   self.power  )
								if not isinstance(self.power    ,   (float, int)  ):
									raise ValueError('Invalid power'  )

								_snake_case   						=		state_dict.get('shadow_params'    ,   lowerCAmelCase_  )
								if shadow_params is not None:
									_snake_case   						=		shadow_params
									if not isinstance(self.shadow_params    ,   lowerCAmelCase_  ):
										raise ValueError('shadow_params must be a list'  )
									if not all(isinstance(lowerCAmelCase_    ,   torch.Tensor  ) for p in self.shadow_params  ):
										raise ValueError('shadow_params must all be Tensors'  )



'''simple docstring'''







import subprocess
import sys

from transformers import BertConfig, BertModel, BertTokenizer, pipeline
from transformers.testing_utils import TestCasePlus, require_torch


class 						__UpperCAmelCase      (					_lowerCamelCase						):


							@require_torch
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

								_snake_case   						=		'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

								_snake_case   						=		'\nimport socket\ndef offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled, we shouldn\'t access internet")\nsocket.socket = offline_socket\n        '

								# Force fetching the files so that we can use the cache
								_snake_case   						=		'hf-internal-testing/tiny-random-bert'
								BertConfig.from_pretrained(lowerCAmelCase_  )
								BertModel.from_pretrained(lowerCAmelCase_  )
								BertTokenizer.from_pretrained(lowerCAmelCase_  )
								pipeline(task='fill-mask'    ,   model=lowerCAmelCase_  )

								# baseline - just load from_pretrained with normal network
								_snake_case   						=		[sys.executable, '-c', '\n'.join([load, run, mock]  )]

								# should succeed
								_snake_case   						=		self.get_env()
								# should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
								_snake_case   						=		'1'
								_snake_case   						=		subprocess.run(lowerCAmelCase_    ,   env=lowerCAmelCase_    ,   check=lowerCAmelCase_    ,   capture_output=lowerCAmelCase_  )
								self.assertEqual(result.returncode    ,   0    ,   result.stderr  )
								self.assertIn('success'    ,   result.stdout.decode()  )


							@require_torch
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

								_snake_case   						=		'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

								_snake_case   						=		'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet")\nsocket.socket = offline_socket\n        '

								# Force fetching the files so that we can use the cache
								_snake_case   						=		'hf-internal-testing/tiny-random-bert'
								BertConfig.from_pretrained(lowerCAmelCase_  )
								BertModel.from_pretrained(lowerCAmelCase_  )
								BertTokenizer.from_pretrained(lowerCAmelCase_  )
								pipeline(task='fill-mask'    ,   model=lowerCAmelCase_  )

								# baseline - just load from_pretrained with normal network
								_snake_case   						=		[sys.executable, '-c', '\n'.join([load, run, mock]  )]

								# should succeed
								_snake_case   						=		self.get_env()
								_snake_case   						=		subprocess.run(lowerCAmelCase_    ,   env=lowerCAmelCase_    ,   check=lowerCAmelCase_    ,   capture_output=lowerCAmelCase_  )
								self.assertEqual(result.returncode    ,   0    ,   result.stderr  )
								self.assertIn('success'    ,   result.stdout.decode()  )


							@require_torch
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		'\nfrom transformers import BertConfig, BertModel, BertTokenizer\n        '

								_snake_case   						=		'\nmname = "hf-internal-testing/tiny-random-bert-sharded"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nprint("success")\n        '

								_snake_case   						=		'\nimport socket\ndef offline_socket(*args, **kwargs): raise ValueError("Offline mode is enabled")\nsocket.socket = offline_socket\n        '

								# baseline - just load from_pretrained with normal network
								_snake_case   						=		[sys.executable, '-c', '\n'.join([load, run]  )]

								# should succeed
								_snake_case   						=		self.get_env()
								_snake_case   						=		subprocess.run(lowerCAmelCase_    ,   env=lowerCAmelCase_    ,   check=lowerCAmelCase_    ,   capture_output=lowerCAmelCase_  )
								self.assertEqual(result.returncode    ,   0    ,   result.stderr  )
								self.assertIn('success'    ,   result.stdout.decode()  )

								# next emulate no network
								_snake_case   						=		[sys.executable, '-c', '\n'.join([load, mock, run]  )]

								# Doesn't fail anymore since the model is in the cache due to other tests, so commenting this.
								# env["TRANSFORMERS_OFFLINE"] = "0"
								# result = subprocess.run(cmd, env=env, check=False, capture_output=True)
								# self.assertEqual(result.returncode, 1, result.stderr)

								# should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
								_snake_case   						=		'1'
								_snake_case   						=		subprocess.run(lowerCAmelCase_    ,   env=lowerCAmelCase_    ,   check=lowerCAmelCase_    ,   capture_output=lowerCAmelCase_  )
								self.assertEqual(result.returncode    ,   0    ,   result.stderr  )
								self.assertIn('success'    ,   result.stdout.decode()  )


							@require_torch
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		'\nfrom transformers import pipeline\n        '
								_snake_case   						=		'\nmname = "hf-internal-testing/tiny-random-bert"\npipe = pipeline(model=mname)\n        '

								_snake_case   						=		'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Offline mode is enabled")\nsocket.socket = offline_socket\n        '
								_snake_case   						=		self.get_env()
								_snake_case   						=		'1'
								_snake_case   						=		[sys.executable, '-c', '\n'.join([load, mock, run]  )]
								_snake_case   						=		subprocess.run(lowerCAmelCase_    ,   env=lowerCAmelCase_    ,   check=lowerCAmelCase_    ,   capture_output=lowerCAmelCase_  )
								self.assertEqual(result.returncode    ,   1    ,   result.stderr  )
								self.assertIn(
								    'You cannot infer task automatically within `pipeline` when using offline mode'    ,   result.stderr.decode().replace('\n'    ,   ''  )    ,   )




							@require_torch
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		'\nfrom transformers import AutoModel\n        '
								_snake_case   						=		'\nmname = "hf-internal-testing/test_dynamic_model"\nAutoModel.from_pretrained(mname, trust_remote_code=True)\nprint("success")\n        '

								# baseline - just load from_pretrained with normal network
								_snake_case   						=		[sys.executable, '-c', '\n'.join([load, run]  )]

								# should succeed
								_snake_case   						=		self.get_env()
								_snake_case   						=		subprocess.run(lowerCAmelCase_    ,   env=lowerCAmelCase_    ,   check=lowerCAmelCase_    ,   capture_output=lowerCAmelCase_  )
								self.assertEqual(result.returncode    ,   0    ,   result.stderr  )
								self.assertIn('success'    ,   result.stdout.decode()  )

								# should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
								_snake_case   						=		'1'
								_snake_case   						=		subprocess.run(lowerCAmelCase_    ,   env=lowerCAmelCase_    ,   check=lowerCAmelCase_    ,   capture_output=lowerCAmelCase_  )
								self.assertEqual(result.returncode    ,   0    ,   result.stderr  )
								self.assertIn('success'    ,   result.stdout.decode()  )

