



"""simple docstring"""







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


lowercase__		:				Any          =	logging.get_logger(__name__)

# General docstring
lowercase__		:				Dict          =	"""RegNetConfig"""

# Base docstring
lowercase__		:				Any          =	"""facebook/regnet-y-040"""
lowercase__		:				Dict          =	[1, 1_0_8_8, 7, 7]

# Image classification docstring
lowercase__		:				List[Any]          =	"""facebook/regnet-y-040"""
lowercase__		:				str          =	"""tabby, tabby cat"""

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







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Optional[Any]    ,			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"    ,			):
        super().__init__()
        lowerCAmelCase_   :				int								= 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_   :				List[Any]								= nn.BatchNormad(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Optional[int]								= ACTaFN[activation] if activation is not None else nn.Identity()





      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        lowerCAmelCase_   :				Dict								= self.convolution(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Tuple								= self.normalization(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				str								= self.activation(SCREAMING_SNAKE_CASE_       )
        return hidden_state







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      RegNetConfig       ):
        super().__init__()
        lowerCAmelCase_   :				int								= RegNetConvLayer(
            config.num_channels    ,			config.embedding_size    ,			kernel_size=3    ,			stride=2    ,			activation=config.hidden_act       )
        lowerCAmelCase_   :				Dict								= config.num_channels





      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        lowerCAmelCase_   :				Optional[Any]								= 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_   :				Optional[Any]								= self.embedder(SCREAMING_SNAKE_CASE_       )
        return hidden_state







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int = 2       ):
        super().__init__()
        lowerCAmelCase_   :				Union[str, Any]								= nn.Convad(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			kernel_size=1    ,			stride=SCREAMING_SNAKE_CASE_    ,			bias=SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Optional[int]								= nn.BatchNormad(SCREAMING_SNAKE_CASE_       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      Tensor       ):
        lowerCAmelCase_   :				int								= self.convolution(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				str								= self.normalization(SCREAMING_SNAKE_CASE_       )
        return hidden_state







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int       ):
        super().__init__()

        lowerCAmelCase_   :				Dict								= nn.AdaptiveAvgPoolad((1, 1)       )
        lowerCAmelCase_   :				Optional[Any]								= 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      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      int       ):
        # b c h w -> b c 1 1
        lowerCAmelCase_   :				Optional[Any]								= self.pooler(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Union[str, Any]								= self.attention(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Tuple								= hidden_state * attention
        return hidden_state







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      RegNetConfig    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int = 1       ):
        super().__init__()
        lowerCAmelCase_   :				int								= in_channels != out_channels or stride != 1
        lowerCAmelCase_   :				int								= max(1    ,			out_channels // config.groups_width       )
        lowerCAmelCase_   :				Dict								= (
            RegNetShortCut(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			stride=SCREAMING_SNAKE_CASE_       ) if should_apply_shortcut else nn.Identity()
        )
        lowerCAmelCase_   :				Union[str, Any]								= 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_   :				Dict								= ACTaFN[config.hidden_act]





      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      Dict       ):
        lowerCAmelCase_   :				Union[str, Any]								= hidden_state
        lowerCAmelCase_   :				str								= self.layer(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Tuple								= self.shortcut(SCREAMING_SNAKE_CASE_       )
        hidden_state += residual
        lowerCAmelCase_   :				Optional[Any]								= self.activation(SCREAMING_SNAKE_CASE_       )
        return hidden_state







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      RegNetConfig    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int = 1       ):
        super().__init__()
        lowerCAmelCase_   :				str								= in_channels != out_channels or stride != 1
        lowerCAmelCase_   :				int								= max(1    ,			out_channels // config.groups_width       )
        lowerCAmelCase_   :				Dict								= (
            RegNetShortCut(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			stride=SCREAMING_SNAKE_CASE_       ) if should_apply_shortcut else nn.Identity()
        )
        lowerCAmelCase_   :				Optional[Any]								= 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_   :				Union[str, Any]								= ACTaFN[config.hidden_act]





      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        lowerCAmelCase_   :				List[str]								= hidden_state
        lowerCAmelCase_   :				int								= self.layer(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				str								= self.shortcut(SCREAMING_SNAKE_CASE_       )
        hidden_state += residual
        lowerCAmelCase_   :				Tuple								= self.activation(SCREAMING_SNAKE_CASE_       )
        return hidden_state







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      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    ,			):
        super().__init__()

        lowerCAmelCase_   :				Dict								= RegNetXLayer if config.layer_type == 'x' else RegNetYLayer

        lowerCAmelCase_   :				Optional[Any]								= 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      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      Any       ):
        lowerCAmelCase_   :				Tuple								= self.layers(SCREAMING_SNAKE_CASE_       )
        return hidden_state







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      RegNetConfig       ):
        super().__init__()
        lowerCAmelCase_   :				Tuple								= 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_   :				Union[str, Any]								= 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      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Tensor    ,			SCREAMING_SNAKE_CASE_				:      bool = False    ,			SCREAMING_SNAKE_CASE_				:      bool = True       ):
        lowerCAmelCase_   :				str								= () if output_hidden_states else None

        for stage_module in self.stages:
          if output_hidden_states:
            lowerCAmelCase_   :				Optional[Any]								= hidden_states + (hidden_state,)

          lowerCAmelCase_   :				Union[str, Any]								= stage_module(SCREAMING_SNAKE_CASE_       )

        if output_hidden_states:
          lowerCAmelCase_   :				Union[str, Any]								= 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       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     RegNetConfig
      _SCREAMING_SNAKE_CASE  				=     """regnet"""
      _SCREAMING_SNAKE_CASE  				=     """pixel_values"""
      _SCREAMING_SNAKE_CASE  				=     True



      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        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      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=False       ):
        if isinstance(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ):
          lowerCAmelCase_   :				Tuple								= value




lowercase__		:				Optional[Any]          =	R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

lowercase__		:				int          =	R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
"""







@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       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def __init__(      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Tuple       ):
        super().__init__(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Optional[int]								= config
        lowerCAmelCase_   :				Optional[Any]								= RegNetEmbeddings(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Any								= RegNetEncoder(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Union[str, Any]								= 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      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Tensor    ,			SCREAMING_SNAKE_CASE_				:      Optional[bool] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[bool] = None       ):
        lowerCAmelCase_   :				Optional[int]								= (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowerCAmelCase_   :				str								= return_dict if return_dict is not None else self.config.use_return_dict

        lowerCAmelCase_   :				Optional[Any]								= self.embedder(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Optional[Any]								= self.encoder(
            SCREAMING_SNAKE_CASE_    ,			output_hidden_states=SCREAMING_SNAKE_CASE_    ,			return_dict=SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Dict								= encoder_outputs[0]

        lowerCAmelCase_   :				int								= 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       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def __init__(      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        super().__init__(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Tuple								= config.num_labels
        lowerCAmelCase_   :				Any								= RegNetModel(SCREAMING_SNAKE_CASE_       )
        # classification head
        lowerCAmelCase_   :				Optional[int]								= 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      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]    ,			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    ,			):
        lowerCAmelCase_   :				Optional[Any]								= return_dict if return_dict is not None else self.config.use_return_dict

        lowerCAmelCase_   :				str								= self.regnet(SCREAMING_SNAKE_CASE_    ,			output_hidden_states=SCREAMING_SNAKE_CASE_    ,			return_dict=SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				str								= outputs.pooler_output if return_dict else outputs[1]

        lowerCAmelCase_   :				Union[str, Any]								= self.classifier(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Dict								= None

        if labels is not None:
          if self.config.problem_type is None:
            if self.num_labels == 1:
              lowerCAmelCase_   :				Dict								= 'regression'
            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
              lowerCAmelCase_   :				Optional[Any]								= 'single_label_classification'
            else:
              lowerCAmelCase_   :				Optional[Any]								= 'multi_label_classification'
          if self.config.problem_type == "regression":
            lowerCAmelCase_   :				str								= MSELoss()
            if self.num_labels == 1:
              lowerCAmelCase_   :				str								= loss_fct(logits.squeeze()    ,			labels.squeeze()       )
            else:
              lowerCAmelCase_   :				int								= loss_fct(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
          elif self.config.problem_type == "single_label_classification":
            lowerCAmelCase_   :				str								= CrossEntropyLoss()
            lowerCAmelCase_   :				Tuple								= loss_fct(logits.view(-1    ,			self.num_labels       )    ,			labels.view(-1       )       )
          elif self.config.problem_type == "multi_label_classification":
            lowerCAmelCase_   :				Union[str, Any]								= BCEWithLogitsLoss()
            lowerCAmelCase_   :				Optional[Any]								= loss_fct(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )

        if not return_dict:
          lowerCAmelCase_   :				Optional[int]								= (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       )










"""simple docstring"""







import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

lowercase__		:				List[str]          =	logging.getLogger()

def 				UpperCamelCase_    (    lowerCAmelCase__      :				str						)						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				Dict								= {}
  lowerCAmelCase_   :				Union[str, Any]								= os.path.join(lowerCAmelCase__    ,   'all_results.json'						)
  if os.path.exists(lowerCAmelCase__						):
    with open(lowerCAmelCase__    ,   'r'						) as f:
      lowerCAmelCase_   :				Any								= json.load(lowerCAmelCase__						)
  else:
    raise ValueError(f"can't find {path}"						)
  return results


lowercase__		:				Optional[int]          =	logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







@require_torch_tpu
class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      Any       ):
        import xla_spawn

        lowerCAmelCase_   :				Union[str, Any]								= self.get_auto_remove_tmp_dir()
        lowerCAmelCase_   :				Optional[Any]								= F"\n            ./examples/pytorch/text-classification/run_glue.py\n            --num_cores=8\n            ./examples/pytorch/text-classification/run_glue.py\n            --model_name_or_path distilbert-base-uncased\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv\n            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv\n            --do_train\n            --do_eval\n            --debug tpu_metrics_debug\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --learning_rate=1e-4\n            --max_steps=10\n            --warmup_steps=2\n            --seed=42\n            --max_seq_length=128\n            ".split()

        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          lowerCAmelCase_   :				Optional[Any]								= time()
          xla_spawn.main()
          lowerCAmelCase_   :				Any								= time()

          lowerCAmelCase_   :				Dict								= get_results(SCREAMING_SNAKE_CASE_       )
          self.assertGreaterEqual(result['eval_accuracy']    ,			0.75       )

          # Assert that the script takes less than 500 seconds to make sure it doesn't hang.
          self.assertLess(end - start    ,			5_0_0       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        import xla_spawn

        lowerCAmelCase_   :				Dict								= '\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            '.split()
        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          xla_spawn.main()





