
import comet  # From: unbabel-comet
import torch

import datasets


lowerCamelCase__    :		Optional[int]          =   datasets.logging.get_logger(__name__)

lowerCamelCase__    :		Optional[Any]          =   """\
@inproceedings{rei-EtAl:2020:WMT,
   author    = {Rei, Ricardo  and  Stewart, Craig  and  Farinha, Ana C  and  Lavie, Alon},
   title     = {Unbabel's Participation in the WMT20 Metrics Shared Task},
   booktitle      = {Proceedings of the Fifth Conference on Machine Translation},
   month          = {November},
   year           = {2020},
   address        = {Online},
   publisher      = {Association for Computational Linguistics},
   pages     = {909--918},
}
@inproceedings{rei-etal-2020-comet,
   title = \"{COMET}: A Neural Framework for {MT} Evaluation\",
   author = \"Rei, Ricardo  and
      Stewart, Craig  and
      Farinha, Ana C  and
      Lavie, Alon\",
   booktitle = \"Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)\",
   month = nov,
   year = \"2020\",
   address = \"Online\",
   publisher = \"Association for Computational Linguistics\",
   url = \"https://www.aclweb.org/anthology/2020.emnlp-main.213\",
   pages = \"2685--2702\",
}
"""

lowerCamelCase__    :		Tuple          =   """\
Crosslingual Optimized Metric for Evaluation of Translation (COMET) is an open-source framework used to train Machine Translation metrics that achieve high levels of correlation with different types of human judgments (HTER, DA's or MQM).
With the release of the framework the authors also released fully trained models that were used to compete in the WMT20 Metrics Shared Task achieving SOTA in that years competition.

See the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.
"""

lowerCamelCase__    :		Optional[int]          =   """
COMET score.

Args:

`sources` (list of str): Source sentences
`predictions` (list of str): candidate translations
`references` (list of str): reference translations
`cuda` (bool): If set to True, runs COMET using GPU
`show_progress` (bool): Shows progress
`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.

Returns:
    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.
    `scores`: List of scores.

Examples:

    >>> comet_metric = datasets.load_metric('comet')
    >>> # comet_metric = load_metric('comet', 'wmt20-comet-da')  # you can also choose which model to use
    >>> source = [\"Dem Feuer konnte Einhalt geboten werden\", \"Schulen und Kindergärten wurden eröffnet.\"]
    >>> hypothesis = [\"The fire could be stopped\", \"Schools and kindergartens were open\"]
    >>> reference = [\"They were able to control the fire.\", \"Schools and kindergartens opened\"]
    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)
    >>> print([round(v, 2) for v in results[\"scores\"]])
    [0.19, 0.92]
"""



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,						_KWARGS_DESCRIPTION							)
class       _snake_case   (	datasets.Metric							):


	def      lowercase__     (					self):


			'''simple docstring'''




			return datasets.MetricInfo(
			    description=_DESCRIPTION					,  citation=_CITATION					,  homepage="""https://unbabel.github.io/COMET/html/index.html"""					,  inputs_description=_KWARGS_DESCRIPTION					,  features=datasets.Features(
			        {
			            """sources""": datasets.Value("""string"""					,  id="""sequence"""),
			            """predictions""": datasets.Value("""string"""					,  id="""sequence"""),
			            """references""": datasets.Value("""string"""					,  id="""sequence"""),
			        })					,  codebase_urls=["""https://github.com/Unbabel/COMET"""]					,  reference_urls=[
			        """https://github.com/Unbabel/COMET""",
			        """https://www.aclweb.org/anthology/2020.emnlp-main.213/""",
			        """http://www.statmt.org/wmt20/pdf/2020.wmt-1.101.pdf6""",
			    ]					,  )


	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			if self.config_name == "default":
					lowercase__  :   str    		=      comet.load_from_checkpoint(comet.download_model("""wmt20-comet-da"""))
			else:
					lowercase__  :   Tuple    		=      comet.load_from_checkpoint(comet.download_model(self.config_name))



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=None					,  SCREAMING_SNAKE_CASE_=False):


			'''simple docstring'''




			if gpus is None:
					lowercase__  :   Any    		=      1 if torch.cuda.is_available() else 0
			lowercase__  :   List[Any]    		=      {"""src""": sources, """mt""": predictions, """ref""": references}
			lowercase__  :   Dict    		=      [dict(zip(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)) for t in zip(*data.values())]
			lowercase__  :   Any    		=      self.scorer.predict(SCREAMING_SNAKE_CASE_					,  gpus=SCREAMING_SNAKE_CASE_					,  progress_bar=SCREAMING_SNAKE_CASE_)
			return {"mean_score": mean_score, "scores": scores}








from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


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

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

# Base docstring
lowerCamelCase__    :		str          =   """facebook/regnet-y-040"""
lowerCamelCase__    :		str          =   [1, 1_0_8_8, 7, 7]

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

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



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = 3					,  SCREAMING_SNAKE_CASE_ = 1					,  SCREAMING_SNAKE_CASE_ = 1					,  SCREAMING_SNAKE_CASE_ = "relu"					,  **SCREAMING_SNAKE_CASE_					,  ):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   # The padding and conv has been verified in
   # https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
   lowercase__  :   str    		=      tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2)
   lowercase__  :   Optional[int]    		=      tf.keras.layers.ConvaD(
       filters=SCREAMING_SNAKE_CASE_					,  kernel_size=SCREAMING_SNAKE_CASE_					,  strides=SCREAMING_SNAKE_CASE_					,  padding="""VALID"""					,  groups=SCREAMING_SNAKE_CASE_					,  use_bias=SCREAMING_SNAKE_CASE_					,  name="""convolution"""					,  )
   lowercase__  :   str    		=      tf.keras.layers.BatchNormalization(epsilon=1E-5					,  momentum=0.9					,  name="""normalization""")
   lowercase__  :   Any    		=      ACTaFN[activation] if activation is not None else tf.identity



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__  :   Tuple    		=      self.convolution(self.padding(SCREAMING_SNAKE_CASE_))
   lowercase__  :   List[Any]    		=      self.normalization(SCREAMING_SNAKE_CASE_)
   lowercase__  :   str    		=      self.activation(SCREAMING_SNAKE_CASE_)
   return hidden_state



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   List[str]    		=      config.num_channels
   lowercase__  :   int    		=      TFRegNetConvLayer(
       out_channels=config.embedding_size					,  kernel_size=3					,  stride=2					,  activation=config.hidden_act					,  name="""embedder"""					,  )



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__  :   List[str]    		=      shape_list(SCREAMING_SNAKE_CASE_)[1]
   if tf.executing_eagerly() and 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.""")

   # When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
   # So change the input format from `NCHW` to `NHWC`.
   # shape = (batch_size, in_height, in_width, in_channels=num_channels)
   lowercase__  :   Optional[int]    		=      tf.transpose(SCREAMING_SNAKE_CASE_					,  perm=(0, 2, 3, 1))
   lowercase__  :   Optional[int]    		=      self.embedder(SCREAMING_SNAKE_CASE_)
   return hidden_state



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = 2					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   Any    		=      tf.keras.layers.ConvaD(
       filters=SCREAMING_SNAKE_CASE_					,  kernel_size=1					,  strides=SCREAMING_SNAKE_CASE_					,  use_bias=SCREAMING_SNAKE_CASE_					,  name="""convolution""")
   lowercase__  :   str    		=      tf.keras.layers.BatchNormalization(epsilon=1E-5					,  momentum=0.9					,  name="""normalization""")



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = False):


   '''simple docstring'''




   return self.normalization(self.convolution(SCREAMING_SNAKE_CASE_)					,  training=SCREAMING_SNAKE_CASE_)



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   str    		=      tf.keras.layers.GlobalAveragePoolingaD(keepdims=SCREAMING_SNAKE_CASE_					,  name="""pooler""")
   lowercase__  :   Tuple    		=      [
       tf.keras.layers.ConvaD(filters=SCREAMING_SNAKE_CASE_					,  kernel_size=1					,  activation="""relu"""					,  name="""attention.0"""),
       tf.keras.layers.ConvaD(filters=SCREAMING_SNAKE_CASE_					,  kernel_size=1					,  activation="""sigmoid"""					,  name="""attention.2"""),
   ]



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__  :   Optional[Any]    		=      self.pooler(SCREAMING_SNAKE_CASE_)
   for layer_module in self.attention:
     lowercase__  :   Optional[Any]    		=      layer_module(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Any    		=      hidden_state * pooled
   return hidden_state



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = 1					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   int    		=      in_channels != out_channels or stride != 1
   lowercase__  :   Any    		=      max(1					,  out_channels // config.groups_width)
   lowercase__  :   List[Any]    		=      (
       TFRegNetShortCut(SCREAMING_SNAKE_CASE_					,  stride=SCREAMING_SNAKE_CASE_					,  name="""shortcut""")
       if should_apply_shortcut
       else tf.keras.layers.Activation("""linear"""					,  name="""shortcut""")
   )
   # `self.layers` instead of `self.layer` because that is a reserved argument.
   lowercase__  :   Optional[int]    		=      [
       TFRegNetConvLayer(SCREAMING_SNAKE_CASE_					,  kernel_size=1					,  activation=config.hidden_act					,  name="""layer.0"""),
       TFRegNetConvLayer(
           SCREAMING_SNAKE_CASE_					,  stride=SCREAMING_SNAKE_CASE_					,  groups=SCREAMING_SNAKE_CASE_					,  activation=config.hidden_act					,  name="""layer.1"""),
       TFRegNetConvLayer(SCREAMING_SNAKE_CASE_					,  kernel_size=1					,  activation=SCREAMING_SNAKE_CASE_					,  name="""layer.2"""),
   ]
   lowercase__  :   Tuple    		=      ACTaFN[config.hidden_act]



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__  :   Tuple    		=      hidden_state
   for layer_module in self.layers:
     lowercase__  :   Any    		=      layer_module(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Union[str, Any]    		=      self.shortcut(SCREAMING_SNAKE_CASE_)
   hidden_state += residual
   lowercase__  :   Any    		=      self.activation(SCREAMING_SNAKE_CASE_)
   return hidden_state



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = 1					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   Any    		=      in_channels != out_channels or stride != 1
   lowercase__  :   str    		=      max(1					,  out_channels // config.groups_width)
   lowercase__  :   int    		=      (
       TFRegNetShortCut(SCREAMING_SNAKE_CASE_					,  stride=SCREAMING_SNAKE_CASE_					,  name="""shortcut""")
       if should_apply_shortcut
       else tf.keras.layers.Activation("""linear"""					,  name="""shortcut""")
   )
   lowercase__  :   Any    		=      [
       TFRegNetConvLayer(SCREAMING_SNAKE_CASE_					,  kernel_size=1					,  activation=config.hidden_act					,  name="""layer.0"""),
       TFRegNetConvLayer(
           SCREAMING_SNAKE_CASE_					,  stride=SCREAMING_SNAKE_CASE_					,  groups=SCREAMING_SNAKE_CASE_					,  activation=config.hidden_act					,  name="""layer.1"""),
       TFRegNetSELayer(SCREAMING_SNAKE_CASE_					,  reduced_channels=int(round(in_channels / 4))					,  name="""layer.2"""),
       TFRegNetConvLayer(SCREAMING_SNAKE_CASE_					,  kernel_size=1					,  activation=SCREAMING_SNAKE_CASE_					,  name="""layer.3"""),
   ]
   lowercase__  :   Union[str, Any]    		=      ACTaFN[config.hidden_act]



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__  :   Union[str, Any]    		=      hidden_state
   for layer_module in self.layers:
     lowercase__  :   Optional[Any]    		=      layer_module(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Optional[int]    		=      self.shortcut(SCREAMING_SNAKE_CASE_)
   hidden_state += residual
   lowercase__  :   Optional[int]    		=      self.activation(SCREAMING_SNAKE_CASE_)
   return hidden_state



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = 2					,  SCREAMING_SNAKE_CASE_ = 2					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)

   lowercase__  :   Any    		=      TFRegNetXLayer if config.layer_type == """x""" else TFRegNetYLayer
   lowercase__  :   Dict    		=      [
       # 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_					,  name="""layers.0"""),
       *[layer(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  name=f'layers.{i+1}') for i in range(depth - 1)],
   ]



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   for layer_module in self.layers:
     lowercase__  :   Tuple    		=      layer_module(SCREAMING_SNAKE_CASE_)
   return hidden_state



class       _snake_case   (	tf.keras.layers.Layer							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   Tuple    		=      []
   # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
   self.stages.append(
       TFRegNetStage(
           SCREAMING_SNAKE_CASE_					,  config.embedding_size					,  config.hidden_sizes[0]					,  stride=2 if config.downsample_in_first_stage else 1					,  depth=config.depths[0]					,  name="""stages.0"""					,  ))
   lowercase__  :   List[Any]    		=      zip(config.hidden_sizes					,  config.hidden_sizes[1:])
   for i, ((in_channels, out_channels), depth) in enumerate(zip(SCREAMING_SNAKE_CASE_					,  config.depths[1:])):
     self.stages.append(TFRegNetStage(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  depth=SCREAMING_SNAKE_CASE_					,  name=f'stages.{i+1}'))



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = False					,  SCREAMING_SNAKE_CASE_ = True):


   '''simple docstring'''




   lowercase__  :   Any    		=      () if output_hidden_states else None

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

     lowercase__  :   str    		=      stage_module(SCREAMING_SNAKE_CASE_)

   if output_hidden_states:
     lowercase__  :   Optional[int]    		=      hidden_states + (hidden_state,)

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

   return TFBaseModelOutputWithNoAttention(last_hidden_state=SCREAMING_SNAKE_CASE_					,  hidden_states=SCREAMING_SNAKE_CASE_)



@keras_serializable
class       _snake_case   (	tf.keras.layers.Layer							):
 __lowerCAmelCase  :     List[Any]        =			RegNetConfig


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   Dict    		=      config
   lowercase__  :   List[Any]    		=      TFRegNetEmbeddings(SCREAMING_SNAKE_CASE_					,  name="""embedder""")
   lowercase__  :   Any    		=      TFRegNetEncoder(SCREAMING_SNAKE_CASE_					,  name="""encoder""")
   lowercase__  :   Optional[int]    		=      tf.keras.layers.GlobalAveragePoolingaD(keepdims=SCREAMING_SNAKE_CASE_					,  name="""pooler""")



 @unpack_inputs
 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = False					,  ):


   '''simple docstring'''




   lowercase__  :   str    		=      (
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   lowercase__  :   str    		=      return_dict if return_dict is not None else self.config.use_return_dict

   lowercase__  :   Any    		=      self.embedder(SCREAMING_SNAKE_CASE_					,  training=SCREAMING_SNAKE_CASE_)

   lowercase__  :   Any    		=      self.encoder(
       SCREAMING_SNAKE_CASE_					,  output_hidden_states=SCREAMING_SNAKE_CASE_					,  return_dict=SCREAMING_SNAKE_CASE_					,  training=SCREAMING_SNAKE_CASE_)

   lowercase__  :   Optional[int]    		=      encoder_outputs[0]
   lowercase__  :   Tuple    		=      self.pooler(SCREAMING_SNAKE_CASE_)

   # Change to NCHW output format have uniformity in the modules
   lowercase__  :   Dict    		=      tf.transpose(SCREAMING_SNAKE_CASE_					,  perm=(0, 3, 1, 2))
   lowercase__  :   Optional[Any]    		=      tf.transpose(SCREAMING_SNAKE_CASE_					,  perm=(0, 3, 1, 2))

   # Change the other hidden state outputs to NCHW as well
   if output_hidden_states:
     lowercase__  :   Optional[int]    		=      tuple([tf.transpose(SCREAMING_SNAKE_CASE_					,  perm=(0, 3, 1, 2)) for h in encoder_outputs[1]])

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

   return TFBaseModelOutputWithPoolingAndNoAttention(
       last_hidden_state=SCREAMING_SNAKE_CASE_					,  pooler_output=SCREAMING_SNAKE_CASE_					,  hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states					,  )



class       _snake_case   (	UpperCAmelCase_							):
 __lowerCAmelCase  :     Dict        =			RegNetConfig
 __lowerCAmelCase  :     Any        =			'regnet'
 __lowerCAmelCase  :     List[str]        =			'pixel_values'


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)					,  dtype=tf.floataa)}

lowerCamelCase__    :		str          =   R"""
    Parameters:
    This model is a Tensorflow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and
    behavior.
        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 [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
"""

lowerCamelCase__    :		Dict          =   R"""
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConveNextImageProcessor.__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 [`~utils.ModelOutput`] instead of a plain tuple.
"""



@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.'  ,						UpperCAmelCase_  ,						)
class       _snake_case   (	UpperCAmelCase_							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(SCREAMING_SNAKE_CASE_					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_)
   lowercase__  :   Any    		=      TFRegNetMainLayer(SCREAMING_SNAKE_CASE_					,  name="""regnet""")



 @unpack_inputs
 @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      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_=False					,  ):


   '''simple docstring'''




   lowercase__  :   int    		=      (
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   lowercase__  :   List[str]    		=      return_dict if return_dict is not None else self.config.use_return_dict

   lowercase__  :   Dict    		=      self.regnet(
       pixel_values=SCREAMING_SNAKE_CASE_					,  output_hidden_states=SCREAMING_SNAKE_CASE_					,  return_dict=SCREAMING_SNAKE_CASE_					,  training=SCREAMING_SNAKE_CASE_					,  )
   if not return_dict:
     return (outputs[0],) + outputs[1:]

   return TFBaseModelOutputWithPoolingAndNoAttention(
       last_hidden_state=outputs.last_hidden_state					,  pooler_output=outputs.pooler_output					,  hidden_states=outputs.hidden_states					,  )



@add_start_docstrings(
    '\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '  ,						UpperCAmelCase_  ,						)
class       _snake_case   (	UpperCAmelCase_  ,						UpperCAmelCase_							):


 def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(SCREAMING_SNAKE_CASE_					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_)
   lowercase__  :   int    		=      config.num_labels
   lowercase__  :   Optional[int]    		=      TFRegNetMainLayer(SCREAMING_SNAKE_CASE_					,  name="""regnet""")
   # classification head
   lowercase__  :   Any    		=      [
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(config.num_labels					,  name="""classifier.1""") if config.num_labels > 0 else tf.identity,
   ]



 @unpack_inputs
 @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      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_=False					,  ):


   '''simple docstring'''




   lowercase__  :   List[Any]    		=      (
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   lowercase__  :   Optional[Any]    		=      return_dict if return_dict is not None else self.config.use_return_dict

   lowercase__  :   List[str]    		=      self.regnet(
       SCREAMING_SNAKE_CASE_					,  output_hidden_states=SCREAMING_SNAKE_CASE_					,  return_dict=SCREAMING_SNAKE_CASE_					,  training=SCREAMING_SNAKE_CASE_)

   lowercase__  :   Dict    		=      outputs.pooler_output if return_dict else outputs[1]

   lowercase__  :   Optional[int]    		=      self.classifier[0](SCREAMING_SNAKE_CASE_)
   lowercase__  :   List[str]    		=      self.classifier[1](SCREAMING_SNAKE_CASE_)

   lowercase__  :   str    		=      None if labels is None else self.hf_compute_loss(labels=SCREAMING_SNAKE_CASE_					,  logits=SCREAMING_SNAKE_CASE_)

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

   return TFSequenceClassifierOutput(loss=SCREAMING_SNAKE_CASE_					,  logits=SCREAMING_SNAKE_CASE_					,  hidden_states=outputs.hidden_states)






