
from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


SCREAMING_SNAKE_CASE__   :			Union[str, Any]												=    logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE__   :			List[str]												=    "RegNetConfig"

# Base docstring
SCREAMING_SNAKE_CASE__   :			Tuple												=    "facebook/regnet-y-040"
SCREAMING_SNAKE_CASE__   :			List[Any]												=    [1, 1_088, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE__   :			Optional[Any]												=    "facebook/regnet-y-040"
SCREAMING_SNAKE_CASE__   :			int												=    "tabby, tabby cat"

SCREAMING_SNAKE_CASE__   :			Tuple												=    [
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]

class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int = 3							,			SCREAMING_SNAKE_CASE__		:							int = 1							,			SCREAMING_SNAKE_CASE__		:							int = 1							,			SCREAMING_SNAKE_CASE__		:							Optional[str] = "relu"							,			)     ->							Tuple:
						super().__init__()
						__lowerCamelCase    					=	nn.Convad(
						    SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=SCREAMING_SNAKE_CASE__							,			stride=SCREAMING_SNAKE_CASE__							,			padding=kernel_size // 2							,			groups=SCREAMING_SNAKE_CASE__							,			bias=SCREAMING_SNAKE_CASE__							,			)
						__lowerCamelCase    					=	nn.BatchNormad(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	ACTaFN[activation] if activation is not None else nn.Identity()







		def  __A       (		self		:							int							,			SCREAMING_SNAKE_CASE__		:							int  )     ->							Any:
						__lowerCamelCase    					=	self.convolution(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.normalization(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.activation(SCREAMING_SNAKE_CASE__  )
						return hidden_state





class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							int							,			SCREAMING_SNAKE_CASE__		:							RegNetConfig  )     ->							str:
						super().__init__()
						__lowerCamelCase    					=	RegNetConvLayer(
						    config.num_channels							,			config.embedding_size							,			kernel_size=3							,			stride=2							,			activation=config.hidden_act  )
						__lowerCamelCase    					=	config.num_channels







		def  __A       (		self		:							Optional[int]							,			SCREAMING_SNAKE_CASE__		:							List[str]  )     ->							Optional[int]:
						__lowerCamelCase    					=	pixel_values.shape[1]
						if num_channels != self.num_channels:
										raise ValueError(
										    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.'''  )
						__lowerCamelCase    					=	self.embedder(SCREAMING_SNAKE_CASE__  )
						return hidden_state





class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							Union[str, Any]							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int = 2  )     ->							Dict:
						super().__init__()
						__lowerCamelCase    					=	nn.Convad(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=1							,			stride=SCREAMING_SNAKE_CASE__							,			bias=SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	nn.BatchNormad(SCREAMING_SNAKE_CASE__  )







		def  __A       (		self		:							Union[str, Any]							,			SCREAMING_SNAKE_CASE__		:							Tensor  )     ->							Tensor:
						__lowerCamelCase    					=	self.convolution(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.normalization(SCREAMING_SNAKE_CASE__  )
						return hidden_state





class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int  )     ->							Tuple:
						super().__init__()

						__lowerCamelCase    					=	nn.AdaptiveAvgPoolad((1, 1)  )
						__lowerCamelCase    					=	nn.Sequential(
						    nn.Convad(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=1  )							,			nn.ReLU()							,			nn.Convad(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=1  )							,			nn.Sigmoid()							,			)







		def  __A       (		self		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							int  )     ->							Union[str, Any]:
						# b c h w -> b c 1 1
						__lowerCamelCase    					=	self.pooler(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.attention(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	hidden_state * attention
						return hidden_state





class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							List[Any]							,			SCREAMING_SNAKE_CASE__		:							RegNetConfig							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int = 1  )     ->							Dict:
						super().__init__()
						__lowerCamelCase    					=	in_channels != out_channels or stride != 1
						__lowerCamelCase    					=	max(1							,			out_channels // config.groups_width  )
						__lowerCamelCase    					=	(
						    RegNetShortCut(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			stride=SCREAMING_SNAKE_CASE__  ) if should_apply_shortcut else nn.Identity()
						)
						__lowerCamelCase    					=	nn.Sequential(
						    RegNetConvLayer(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=1							,			activation=config.hidden_act  )							,			RegNetConvLayer(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			stride=SCREAMING_SNAKE_CASE__							,			groups=SCREAMING_SNAKE_CASE__							,			activation=config.hidden_act  )							,			RegNetConvLayer(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=1							,			activation=SCREAMING_SNAKE_CASE__  )							,			)
						__lowerCamelCase    					=	ACTaFN[config.hidden_act]







		def  __A       (		self		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Tuple  )     ->							List[Any]:
						__lowerCamelCase    					=	hidden_state
						__lowerCamelCase    					=	self.layer(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.shortcut(SCREAMING_SNAKE_CASE__  )
						hidden_state += residual
						__lowerCamelCase    					=	self.activation(SCREAMING_SNAKE_CASE__  )
						return hidden_state





class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							int							,			SCREAMING_SNAKE_CASE__		:							RegNetConfig							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int = 1  )     ->							List[str]:
						super().__init__()
						__lowerCamelCase    					=	in_channels != out_channels or stride != 1
						__lowerCamelCase    					=	max(1							,			out_channels // config.groups_width  )
						__lowerCamelCase    					=	(
						    RegNetShortCut(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			stride=SCREAMING_SNAKE_CASE__  ) if should_apply_shortcut else nn.Identity()
						)
						__lowerCamelCase    					=	nn.Sequential(
						    RegNetConvLayer(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=1							,			activation=config.hidden_act  )							,			RegNetConvLayer(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			stride=SCREAMING_SNAKE_CASE__							,			groups=SCREAMING_SNAKE_CASE__							,			activation=config.hidden_act  )							,			RegNetSELayer(SCREAMING_SNAKE_CASE__							,			reduced_channels=int(round(in_channels / 4  )  )  )							,			RegNetConvLayer(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			kernel_size=1							,			activation=SCREAMING_SNAKE_CASE__  )							,			)
						__lowerCamelCase    					=	ACTaFN[config.hidden_act]







		def  __A       (		self		:							Dict							,			SCREAMING_SNAKE_CASE__		:							Dict  )     ->							Optional[Any]:
						__lowerCamelCase    					=	hidden_state
						__lowerCamelCase    					=	self.layer(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.shortcut(SCREAMING_SNAKE_CASE__  )
						hidden_state += residual
						__lowerCamelCase    					=	self.activation(SCREAMING_SNAKE_CASE__  )
						return hidden_state





class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							str							,			SCREAMING_SNAKE_CASE__		:							RegNetConfig							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							int = 2							,			SCREAMING_SNAKE_CASE__		:							int = 2							,			)     ->							int:
						super().__init__()

						__lowerCamelCase    					=	RegNetXLayer if config.layer_type == '''x''' else RegNetYLayer

						__lowerCamelCase    					=	nn.Sequential(
						    # downsampling is done in the first layer with stride of 2
						    layer(
						        SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			stride=SCREAMING_SNAKE_CASE__							,			)							,			*[layer(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ) for _ in range(depth - 1  )]							,			)







		def  __A       (		self		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							Optional[int]  )     ->							List[str]:
						__lowerCamelCase    					=	self.layers(SCREAMING_SNAKE_CASE__  )
						return hidden_state





class    lowerCAmelCase__     (	nn.Module			):






		def __init__(		self		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							RegNetConfig  )     ->							int:
						super().__init__()
						__lowerCamelCase    					=	nn.ModuleList([]  )
						# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
						self.stages.append(
						    RegNetStage(
						        SCREAMING_SNAKE_CASE__							,			config.embedding_size							,			config.hidden_sizes[0]							,			stride=2 if config.downsample_in_first_stage else 1							,			depth=config.depths[0]							,			)  )
						__lowerCamelCase    					=	zip(config.hidden_sizes							,			config.hidden_sizes[1:]  )
						for (in_channels, out_channels), depth in zip(SCREAMING_SNAKE_CASE__							,			config.depths[1:]  ):
										self.stages.append(RegNetStage(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			depth=SCREAMING_SNAKE_CASE__  )  )







		def  __A       (		self		:							Dict							,			SCREAMING_SNAKE_CASE__		:							Tensor							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							bool = True  )     ->							BaseModelOutputWithNoAttention:
						__lowerCamelCase    					=	() if output_hidden_states else None

						for stage_module in self.stages:
										if output_hidden_states:
														__lowerCamelCase    					=	hidden_states + (hidden_state,)

										__lowerCamelCase    					=	stage_module(SCREAMING_SNAKE_CASE__  )

						if output_hidden_states:
										__lowerCamelCase    					=	hidden_states + (hidden_state,)

						if not return_dict:
										return tuple(v for v in [hidden_state, hidden_states] if v is not None  )

						return BaseModelOutputWithNoAttention(last_hidden_state=SCREAMING_SNAKE_CASE__							,			hidden_states=SCREAMING_SNAKE_CASE__  )





class    lowerCAmelCase__     (	__lowercase			):
		a__       :						List[Any]        =		RegNetConfig
		a__       :						int        =		"""regnet"""
		a__       :						Optional[Any]        =		"""pixel_values"""
		a__       :						List[Any]        =		True






		def  __A       (		self		:							int							,			SCREAMING_SNAKE_CASE__		:							List[str]  )     ->							str:
						if isinstance(SCREAMING_SNAKE_CASE__							,			nn.Convad  ):
										nn.init.kaiming_normal_(module.weight							,			mode='''fan_out'''							,			nonlinearity='''relu'''  )
						elif isinstance(SCREAMING_SNAKE_CASE__							,			(nn.BatchNormad, nn.GroupNorm)  ):
										nn.init.constant_(module.weight							,			1  )
										nn.init.constant_(module.bias							,			0  )







		def  __A       (		self		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Dict							,			SCREAMING_SNAKE_CASE__		:							Tuple=False  )     ->							List[Any]:
						if isinstance(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ):
										__lowerCamelCase    					=	value





SCREAMING_SNAKE_CASE__   :			str												=    r"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

SCREAMING_SNAKE_CASE__   :			Optional[int]												=    r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.\n"

@add_start_docstrings(
    """The bare RegNet model outputting raw features without any specific head on top."""  ,	__lowercase  ,	)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class    lowerCAmelCase__     (	__lowercase			):






		def __init__(		self		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							Optional[Any]:
						super().__init__(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	config
						__lowerCamelCase    					=	RegNetEmbeddings(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	RegNetEncoder(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	nn.AdaptiveAvgPoolad((1, 1)  )
						# Initialize weights and apply final processing
						self.post_init()







		@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE__  )
		@add_code_sample_docstrings(
		    checkpoint=_CHECKPOINT_FOR_DOC							,			output_type=SCREAMING_SNAKE_CASE__							,			config_class=_CONFIG_FOR_DOC							,			modality='''vision'''							,			expected_output=_EXPECTED_OUTPUT_SHAPE							,			)
		def  __A       (		self		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Tensor							,			SCREAMING_SNAKE_CASE__		:							Optional[bool] = None							,			SCREAMING_SNAKE_CASE__		:							Optional[bool] = None  )     ->							BaseModelOutputWithPoolingAndNoAttention:
						__lowerCamelCase    					=	(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						__lowerCamelCase    					=	return_dict if return_dict is not None else self.config.use_return_dict

						__lowerCamelCase    					=	self.embedder(SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	self.encoder(
						    SCREAMING_SNAKE_CASE__							,			output_hidden_states=SCREAMING_SNAKE_CASE__							,			return_dict=SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	encoder_outputs[0]

						__lowerCamelCase    					=	self.pooler(SCREAMING_SNAKE_CASE__  )

						if not return_dict:
										return (last_hidden_state, pooled_output) + encoder_outputs[1:]

						return BaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=SCREAMING_SNAKE_CASE__							,			pooler_output=SCREAMING_SNAKE_CASE__							,			hidden_states=encoder_outputs.hidden_states							,			)





@add_start_docstrings(
    """
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """  ,	__lowercase  ,	)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class    lowerCAmelCase__     (	__lowercase			):






		def __init__(		self		:							Dict							,			SCREAMING_SNAKE_CASE__		:							List[str]  )     ->							List[Any]:
						super().__init__(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	config.num_labels
						__lowerCamelCase    					=	RegNetModel(SCREAMING_SNAKE_CASE__  )
						# classification head
						__lowerCamelCase    					=	nn.Sequential(
						    nn.Flatten()							,			nn.Linear(config.hidden_sizes[-1]							,			config.num_labels  ) if config.num_labels > 0 else nn.Identity()							,			)
						# initialize weights and apply final processing
						self.post_init()







		@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE__  )
		@add_code_sample_docstrings(
		    checkpoint=_IMAGE_CLASS_CHECKPOINT							,			output_type=SCREAMING_SNAKE_CASE__							,			config_class=_CONFIG_FOR_DOC							,			expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,			)
		def  __A       (		self		:							Dict							,			SCREAMING_SNAKE_CASE__		:							Optional[torch.FloatTensor] = None							,			SCREAMING_SNAKE_CASE__		:							Optional[torch.LongTensor] = None							,			SCREAMING_SNAKE_CASE__		:							Optional[bool] = None							,			SCREAMING_SNAKE_CASE__		:							Optional[bool] = None							,			)     ->							ImageClassifierOutputWithNoAttention:
						__lowerCamelCase    					=	return_dict if return_dict is not None else self.config.use_return_dict

						__lowerCamelCase    					=	self.regnet(SCREAMING_SNAKE_CASE__							,			output_hidden_states=SCREAMING_SNAKE_CASE__							,			return_dict=SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	outputs.pooler_output if return_dict else outputs[1]

						__lowerCamelCase    					=	self.classifier(SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	None

						if labels is not None:
										if self.config.problem_type is None:
														if self.num_labels == 1:
																		__lowerCamelCase    					=	'''regression'''
														elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																		__lowerCamelCase    					=	'''single_label_classification'''
														else:
																		__lowerCamelCase    					=	'''multi_label_classification'''
										if self.config.problem_type == "regression":
														__lowerCamelCase    					=	MSELoss()
														if self.num_labels == 1:
																		__lowerCamelCase    					=	loss_fct(logits.squeeze()							,			labels.squeeze()  )
														else:
																		__lowerCamelCase    					=	loss_fct(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  )
										elif self.config.problem_type == "single_label_classification":
														__lowerCamelCase    					=	CrossEntropyLoss()
														__lowerCamelCase    					=	loss_fct(logits.view(-1							,			self.num_labels  )							,			labels.view(-1  )  )
										elif self.config.problem_type == "multi_label_classification":
														__lowerCamelCase    					=	BCEWithLogitsLoss()
														__lowerCamelCase    					=	loss_fct(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  )

						if not return_dict:
										__lowerCamelCase    					=	(logits,) + outputs[2:]
										return (loss,) + output if loss is not None else output

						return ImageClassifierOutputWithNoAttention(loss=SCREAMING_SNAKE_CASE__							,			logits=SCREAMING_SNAKE_CASE__							,			hidden_states=outputs.hidden_states  )





import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


SCREAMING_SNAKE_CASE__   :			Optional[int]												=    os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(F'{bindir}/../../examples/pytorch/translation'):
					from run_translation import main  # noqa


set_seed(42)
SCREAMING_SNAKE_CASE__   :			Any												=    "sshleifer/student_marian_en_ro_6_1"
SCREAMING_SNAKE_CASE__   :			Tuple												=    "sshleifer/tiny-mbart"

@require_torch
class    lowerCAmelCase__     (	__lowercase			):






		def  __A       (		self		:							List[Any]							,			SCREAMING_SNAKE_CASE__		:							Tuple=False							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=None							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=True							,			SCREAMING_SNAKE_CASE__		:							str=True							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=True							,			SCREAMING_SNAKE_CASE__		:							Optional[int]=True							,			)     ->							Optional[int]:
						__lowerCamelCase    					=	self.run_trainer(
						    eval_steps=1							,			max_len=12							,			model_name=SCREAMING_SNAKE_CASE__							,			num_train_epochs=1							,			distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str=SCREAMING_SNAKE_CASE__							,			predict_with_generate=SCREAMING_SNAKE_CASE__							,			do_train=SCREAMING_SNAKE_CASE__							,			do_eval=SCREAMING_SNAKE_CASE__							,			do_predict=SCREAMING_SNAKE_CASE__							,			)
						__lowerCamelCase    					=	TrainerState.load_from_json(os.path.join(SCREAMING_SNAKE_CASE__							,			'''trainer_state.json'''  )  ).log_history

						if not do_eval:
										return

						__lowerCamelCase    					=	[log for log in logs if '''eval_loss''' in log.keys()]

						__lowerCamelCase    					=	eval_metrics[0]
						if predict_with_generate:
										assert "eval_bleu" in first_step_stats

										__lowerCamelCase    					=	eval_metrics[-1]
										assert isinstance(last_step_stats['''eval_bleu''']							,			SCREAMING_SNAKE_CASE__  )
										assert not math.isnan(float(last_step_stats['''eval_loss''']  )  ), "eval_loss must not be `nan`"






		@require_torch_non_multi_gpu
		def  __A       (		self		:							Optional[int]  )     ->							int:
						self.run_seqaseq_quick()






		@require_torch_multi_gpu
		def  __A       (		self		:							int  )     ->							List[str]:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__  )






		@require_torch_multi_gpu
		def  __A       (		self		:							Optional[Any]  )     ->							Tuple:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Dict  )     ->							Tuple:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp simple'''  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Optional[int]  )     ->							List[str]:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp simple --fp16'''  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Tuple  )     ->							Any:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp zero_dp_2'''							,			predict_with_generate=SCREAMING_SNAKE_CASE__  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Dict  )     ->							Tuple:
						self.run_seqaseq_quick(
						    distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp zero_dp_2 --fp16'''							,			predict_with_generate=SCREAMING_SNAKE_CASE__  )






		@require_apex
		@require_torch_gpu
		def  __A       (		self		:							Union[str, Any]  )     ->							List[str]:
						# XXX: apex breaks the trainer if it's run twice e.g. run_seq2seq.main() from the same
						# program and it breaks other tests that run from the same pytest worker, therefore until this is
						# sorted out it must be run only in an external program, that is distributed=True in this
						# test and only under one or more gpus - if we want cpu will need to make a special test
						#
						# specifically to the problem traced it to self.optimizer.step() - if it's run 2nd time via
						# 2nd main() call it botches the future eval.
						#
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--fp16 --fp16_backend=apex'''  )
						# test 2nd time - was getting eval_loss': nan'
						# to reproduce the problem set distributed=False
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--fp16 --fp16_backend=apex'''  )






		@parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']  )
		@require_torch_multi_gpu
		def  __A       (		self		:							Any							,			SCREAMING_SNAKE_CASE__		:							List[str]  )     ->							Optional[Any]:
						# as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
						__lowerCamelCase    					=	{
						    # test with the default log_level - should be info and thus log info once
						    '''base''': {'''extra_args_str''': '''''', '''n_matches''': 1},
						    # test with low log_level and log_level_replica - should be noisy on all processes
						    # now the info string should appear twice on 2 processes
						    '''low''': {'''extra_args_str''': '''--log_level debug --log_level_replica debug''', '''n_matches''': 2},
						    # test with high log_level and low log_level_replica
						    # now the info string should appear once only on the replica
						    '''high''': {'''extra_args_str''': '''--log_level error --log_level_replica debug''', '''n_matches''': 1},
						    # test with high log_level and log_level_replica - should be quiet on all processes
						    '''mixed''': {'''extra_args_str''': '''--log_level error --log_level_replica error''', '''n_matches''': 0},
						}

						__lowerCamelCase    					=	experiments[experiment_id]
						__lowerCamelCase    					=	{'''distributed''': True, '''predict_with_generate''': False, '''do_eval''': False, '''do_predict''': False}
						__lowerCamelCase    					=	'''Running training'''
						with CaptureStderr() as cl:
										self.run_seqaseq_quick(**SCREAMING_SNAKE_CASE__							,			extra_args_str=data['''extra_args_str''']  )
						__lowerCamelCase    					=	len(re.findall(SCREAMING_SNAKE_CASE__							,			cl.err  )  )
						self.assertEqual(SCREAMING_SNAKE_CASE__							,			data['''n_matches''']  )






		@slow
		def  __A       (		self		:							Any  )     ->							Optional[Any]:
						__lowerCamelCase    					=	self.run_trainer(
						    eval_steps=2							,			max_len=1_28							,			model_name=SCREAMING_SNAKE_CASE__							,			learning_rate=3e-4							,			num_train_epochs=10							,			distributed=SCREAMING_SNAKE_CASE__							,			)

						# Check metrics
						__lowerCamelCase    					=	TrainerState.load_from_json(os.path.join(SCREAMING_SNAKE_CASE__							,			'''trainer_state.json'''  )  ).log_history
						__lowerCamelCase    					=	[log for log in logs if '''eval_loss''' in log.keys()]
						__lowerCamelCase    					=	eval_metrics[0]
						__lowerCamelCase    					=	eval_metrics[-1]

						assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
						assert isinstance(last_step_stats['''eval_bleu''']							,			SCREAMING_SNAKE_CASE__  )

						# test if do_predict saves generations and metrics
						__lowerCamelCase    					=	os.listdir(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	{os.path.basename(SCREAMING_SNAKE_CASE__  ) for p in contents}
						assert "generated_predictions.txt" in contents
						assert "predict_results.json" in contents






		@slow
		@require_bitsandbytes
		def  __A       (		self		:							Optional[int]  )     ->							str:
						from transformers.training_args import OptimizerNames

						def train_and_return_metrics(SCREAMING_SNAKE_CASE__		:							str  ) -> Tuple[int, float]:
										__lowerCamelCase    					=	'''--skip_memory_metrics 0'''

										__lowerCamelCase    					=	self.run_trainer(
										    max_len=1_28							,			model_name=SCREAMING_SNAKE_CASE__							,			learning_rate=3e-4							,			num_train_epochs=1							,			optim=SCREAMING_SNAKE_CASE__							,			distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str=SCREAMING_SNAKE_CASE__							,			do_eval=SCREAMING_SNAKE_CASE__							,			do_predict=SCREAMING_SNAKE_CASE__							,			n_gpus_to_use=1							,			)

										# Check metrics
										__lowerCamelCase    					=	TrainerState.load_from_json(Path(SCREAMING_SNAKE_CASE__							,			'''trainer_state.json'''  )  ).log_history
										__lowerCamelCase    					=	int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**20  )
										__lowerCamelCase    					=	int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**20  )

										__lowerCamelCase    					=	logs[0]['''train_loss''']
										return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value  )
						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	train_and_return_metrics(OptimizerNames.ADAMW_BNB.value  )

						__lowerCamelCase    					=	gpu_alloc_mem_orig - gpu_alloc_mem_bnb

						__lowerCamelCase    					=	gpu_peak_mem_orig + gpu_alloc_mem_orig
						__lowerCamelCase    					=	gpu_peak_mem_bnb + gpu_alloc_mem_bnb
						__lowerCamelCase    					=	gpu_total_mem_orig - gpu_total_mem_bnb

						# sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
						# doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
						# in 2 bytes and the diff in optim memory usage is derived as so:
						#
						# - normal 25*8=~200MB (8 bytes per param)
						# - bnb    25*2= ~50MB (2 bytes per param)
						#
						# Thus we should expect ~150MB total memory saved.
						#
						# Peak memory should be the same - the total should be different by about that same margin
						#
						# After leaving a small margin to accommodate for differences between gpus let's check
						# that we have at least 120MB in savings
						__lowerCamelCase    					=	1_20

						# uncomment the following if this test starts failing - requires py38 for a new print feature
						# gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
						# print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
						# print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
						# print(f"{gpu_alloc_mem_diff=}MB")
						# print(f"{gpu_peak_mem_diff=}MB")
						# print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
						# print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

						self.assertGreater(
						    SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			'''should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'''
						    f''' a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and'''
						    f''' gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB'''							,			)

						self.assertGreater(
						    SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			'''should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'''
						    f''' a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and'''
						    f''' gpu_total_mem_bnb={gpu_total_mem_bnb}MB'''							,			)

						self.assertEqual(
						    SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			f'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''  )







		def  __A       (		self		:							Dict							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							str							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							float = 3e-3							,			SCREAMING_SNAKE_CASE__		:							str = "adafactor"							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							str = None							,			SCREAMING_SNAKE_CASE__		:							int = 0							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							int = None							,			)     ->							List[Any]:
						__lowerCamelCase    					=	self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
						__lowerCamelCase    					=	self.get_auto_remove_tmp_dir()
						__lowerCamelCase    					=	f'''
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(SCREAMING_SNAKE_CASE__  )}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(SCREAMING_SNAKE_CASE__  )}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        '''.split()

						__lowerCamelCase    					=	f'''
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(SCREAMING_SNAKE_CASE__  )}
        '''.split()

						__lowerCamelCase    					=	'''
            --do_predict
        '''.split()

						__lowerCamelCase    					=	[]
						if do_train:
										args += args_train

						if do_eval:
										args += args_eval

						if do_predict:
										args += args_predict

						if predict_with_generate:
										args += "--predict_with_generate".split()

						if do_train:
										if optim == "adafactor":
														args += "--adafactor".split()
										else:
														args += f'''--optim {optim}'''.split()

						if extra_args_str is not None:
										args += extra_args_str.split()

						if distributed:
										if n_gpus_to_use is None:
														__lowerCamelCase    					=	get_gpu_count()
										__lowerCamelCase    					=	get_torch_dist_unique_port()
										__lowerCamelCase    					=	f'''
                -m torch.distributed.run
                --nproc_per_node={n_gpus_to_use}
                --master_port={master_port}
                {self.examples_dir_str}/pytorch/translation/run_translation.py
            '''.split()
										__lowerCamelCase    					=	[sys.executable] + distributed_args + args
										# keep for quick debug
										# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
										execute_subprocess_async(SCREAMING_SNAKE_CASE__							,			env=self.get_env()  )
						else:
										__lowerCamelCase    					=	['''run_translation.py'''] + args
										with patch.object(SCREAMING_SNAKE_CASE__							,			'''argv'''							,			SCREAMING_SNAKE_CASE__  ):
														main()

						return output_dir



