
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      =logging.get_logger(__name__)

# General docstring
lowerCamelCase      ="""RegNetConfig"""

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

# Image classification docstring
lowerCamelCase      ="""facebook/regnet-y-040"""
lowerCamelCase      ="""tabby, tabby cat"""

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




class 				_lowerCamelCase (      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						,						)  ->   str:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   str:




           """simple docstring"""







           UpperCamelCase__    :   Optional[Any]					=				self.convolution(self.padding(_lowerCAmelCase  )  )
           UpperCamelCase__    :   List[str]					=				self.normalization(_lowerCAmelCase  )
           UpperCamelCase__    :   Union[str, Any]					=				self.activation(_lowerCAmelCase  )
           return hidden_state




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






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   Union[str, Any]:




           """simple docstring"""







           super().__init__(**_lowerCAmelCase  )
           UpperCamelCase__    :   List[Any]					=				config.num_channels
           UpperCamelCase__    :   Union[str, Any]					=				TFRegNetConvLayer(
               out_channels=config.embedding_size						,						kernel_size=3						,						stride=2						,						activation=config.hidden_act						,						name='''embedder'''						,						)






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   Optional[int]:




           """simple docstring"""







           UpperCamelCase__    :   int					=				shape_list(_lowerCAmelCase  )[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)
           UpperCamelCase__    :   Optional[int]					=				tf.transpose(_lowerCAmelCase						,						perm=(0, 2, 3, 1)  )
           UpperCamelCase__    :   List[Any]					=				self.embedder(_lowerCAmelCase  )
           return hidden_state




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






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE = 2						,						**__SCREAMING_SNAKE_CASE  )  ->   List[str]:




           """simple docstring"""







           super().__init__(**_lowerCAmelCase  )
           UpperCamelCase__    :   List[str]					=				tf.keras.layers.ConvaD(
               filters=_lowerCAmelCase						,						kernel_size=1						,						strides=_lowerCAmelCase						,						use_bias=_lowerCAmelCase						,						name='''convolution'''  )
           UpperCamelCase__    :   str					=				tf.keras.layers.BatchNormalization(epsilon=1e-5						,						momentum=0.9						,						name='''normalization'''  )






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE = False  )  ->   tf.Tensor:




           """simple docstring"""







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




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






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   Dict:




           """simple docstring"""







           super().__init__(**_lowerCAmelCase  )
           UpperCamelCase__    :   int					=				tf.keras.layers.GlobalAveragePoolingaD(keepdims=_lowerCAmelCase						,						name='''pooler'''  )
           UpperCamelCase__    :   Optional[Any]					=				[
               tf.keras.layers.ConvaD(filters=_lowerCAmelCase						,						kernel_size=1						,						activation='''relu'''						,						name='''attention.0'''  ),
               tf.keras.layers.ConvaD(filters=_lowerCAmelCase						,						kernel_size=1						,						activation='''sigmoid'''						,						name='''attention.2'''  ),
           ]






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   Optional[Any]:




           """simple docstring"""







           UpperCamelCase__    :   Optional[int]					=				self.pooler(_lowerCAmelCase  )
           for layer_module in self.attention:
                 UpperCamelCase__    :   Any					=				layer_module(_lowerCAmelCase  )
           UpperCamelCase__    :   List[str]					=				hidden_state * pooled
           return hidden_state




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






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE = 1						,						**__SCREAMING_SNAKE_CASE  )  ->   List[str]:




           """simple docstring"""







           super().__init__(**_lowerCAmelCase  )
           UpperCamelCase__    :   Optional[int]					=				in_channels != out_channels or stride != 1
           UpperCamelCase__    :   Optional[Any]					=				max(1						,						out_channels // config.groups_width  )
           UpperCamelCase__    :   Any					=				(
               TFRegNetShortCut(_lowerCAmelCase						,						stride=_lowerCAmelCase						,						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.
           UpperCamelCase__    :   int					=				[
               TFRegNetConvLayer(_lowerCAmelCase						,						kernel_size=1						,						activation=config.hidden_act						,						name='''layer.0'''  ),
               TFRegNetConvLayer(
                   _lowerCAmelCase						,						stride=_lowerCAmelCase						,						groups=_lowerCAmelCase						,						activation=config.hidden_act						,						name='''layer.1'''  ),
               TFRegNetConvLayer(_lowerCAmelCase						,						kernel_size=1						,						activation=_lowerCAmelCase						,						name='''layer.2'''  ),
           ]
           UpperCamelCase__    :   str					=				ACTaFN[config.hidden_act]






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   int:




           """simple docstring"""







           UpperCamelCase__    :   Union[str, Any]					=				hidden_state
           for layer_module in self.layers:
                 UpperCamelCase__    :   Optional[Any]					=				layer_module(_lowerCAmelCase  )
           UpperCamelCase__    :   Any					=				self.shortcut(_lowerCAmelCase  )
           hidden_state += residual
           UpperCamelCase__    :   List[str]					=				self.activation(_lowerCAmelCase  )
           return hidden_state




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






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE = 1						,						**__SCREAMING_SNAKE_CASE  )  ->   Tuple:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   Dict:




           """simple docstring"""







           UpperCamelCase__    :   Optional[int]					=				hidden_state
           for layer_module in self.layers:
                 UpperCamelCase__    :   List[Any]					=				layer_module(_lowerCAmelCase  )
           UpperCamelCase__    :   Optional[int]					=				self.shortcut(_lowerCAmelCase  )
           hidden_state += residual
           UpperCamelCase__    :   List[Any]					=				self.activation(_lowerCAmelCase  )
           return hidden_state




class 				_lowerCamelCase (      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  )  ->   Optional[Any]:




           """simple docstring"""







           super().__init__(**_lowerCAmelCase  )

           UpperCamelCase__    :   Any					=				TFRegNetXLayer if config.layer_type == '''x''' else TFRegNetYLayer
           UpperCamelCase__    :   Optional[Any]					=				[
               # downsampling is done in the first layer with stride of 2
               layer(_lowerCAmelCase						,						_lowerCAmelCase						,						_lowerCAmelCase						,						stride=_lowerCAmelCase						,						name='''layers.0'''  ),
               *[layer(_lowerCAmelCase						,						_lowerCAmelCase						,						_lowerCAmelCase						,						name=F'''layers.{i+1}'''  ) for i in range(depth - 1  )],
           ]






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   Optional[int]:




           """simple docstring"""







           for layer_module in self.layers:
                 UpperCamelCase__    :   str					=				layer_module(_lowerCAmelCase  )
           return hidden_state




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






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   Union[str, Any]:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE = False						,						__SCREAMING_SNAKE_CASE = True  )  ->   TFBaseModelOutputWithNoAttention:




           """simple docstring"""







           UpperCamelCase__    :   Dict					=				() if output_hidden_states else None

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

                 UpperCamelCase__    :   List[Any]					=				stage_module(_lowerCAmelCase  )

           if output_hidden_states:
                 UpperCamelCase__    :   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 TFBaseModelOutputWithNoAttention(last_hidden_state=_lowerCAmelCase						,						hidden_states=_lowerCAmelCase  )




@keras_serializable
class 				_lowerCamelCase (      tf.keras.layers.Layer ):
     SCREAMING_SNAKE_CASE_		   =  RegNetConfig






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   Union[str, Any]:




           """simple docstring"""







           super().__init__(**_lowerCAmelCase  )
           UpperCamelCase__    :   Union[str, Any]					=				config
           UpperCamelCase__    :   str					=				TFRegNetEmbeddings(_lowerCAmelCase						,						name='''embedder'''  )
           UpperCamelCase__    :   Any					=				TFRegNetEncoder(_lowerCAmelCase						,						name='''encoder'''  )
           UpperCamelCase__    :   Dict					=				tf.keras.layers.GlobalAveragePoolingaD(keepdims=_lowerCAmelCase						,						name='''pooler'''  )






     @unpack_inputs
     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE = False						,						)  ->   TFBaseModelOutputWithPoolingAndNoAttention:




           """simple docstring"""







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

           UpperCamelCase__    :   Optional[Any]					=				self.embedder(_lowerCAmelCase						,						training=_lowerCAmelCase  )

           UpperCamelCase__    :   Union[str, Any]					=				self.encoder(
               _lowerCAmelCase						,						output_hidden_states=_lowerCAmelCase						,						return_dict=_lowerCAmelCase						,						training=_lowerCAmelCase  )

           UpperCamelCase__    :   Any					=				encoder_outputs[0]
           UpperCamelCase__    :   Tuple					=				self.pooler(_lowerCAmelCase  )

           # Change to NCHW output format have uniformity in the modules
           UpperCamelCase__    :   List[str]					=				tf.transpose(_lowerCAmelCase						,						perm=(0, 3, 1, 2)  )
           UpperCamelCase__    :   str					=				tf.transpose(_lowerCAmelCase						,						perm=(0, 3, 1, 2)  )

           # Change the other hidden state outputs to NCHW as well
           if output_hidden_states:
                 UpperCamelCase__    :   str					=				tuple([tf.transpose(_lowerCAmelCase						,						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=_lowerCAmelCase						,						pooler_output=_lowerCAmelCase						,						hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states						,						)




class 				_lowerCamelCase (      _lowerCAmelCase ):
     SCREAMING_SNAKE_CASE_		   =  RegNetConfig
     SCREAMING_SNAKE_CASE_		   =  'regnet'
     SCREAMING_SNAKE_CASE_		   =  'pixel_values'






     @property
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







           return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_2_4, 2_2_4)						,						dtype=tf.floataa  )}
lowerCamelCase      =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      =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.'''	,							_lowerCAmelCase	,							)
class 				_lowerCamelCase (      _lowerCAmelCase ):






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						*__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   List[Any]:




           """simple docstring"""







           super().__init__(_lowerCAmelCase						,						*_lowerCAmelCase						,						**_lowerCAmelCase  )
           UpperCamelCase__    :   int					=				TFRegNetMainLayer(_lowerCAmelCase						,						name='''regnet'''  )






     @unpack_inputs
     @add_start_docstrings_to_model_forward(_lowerCAmelCase  )
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC						,						output_type=_lowerCAmelCase						,						config_class=_CONFIG_FOR_DOC						,						modality='''vision'''						,						expected_output=_EXPECTED_OUTPUT_SHAPE						,						)
     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE=False						,						)  ->   Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:




           """simple docstring"""







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

           UpperCamelCase__    :   Dict					=				self.regnet(
               pixel_values=_lowerCAmelCase						,						output_hidden_states=_lowerCAmelCase						,						return_dict=_lowerCAmelCase						,						training=_lowerCAmelCase						,						)
           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    '''	,							_lowerCAmelCase	,							)
class 				_lowerCamelCase (      _lowerCAmelCase	,							_lowerCAmelCase ):






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						*__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   int:




           """simple docstring"""







           super().__init__(_lowerCAmelCase						,						*_lowerCAmelCase						,						**_lowerCAmelCase  )
           UpperCamelCase__    :   Any					=				config.num_labels
           UpperCamelCase__    :   List[Any]					=				TFRegNetMainLayer(_lowerCAmelCase						,						name='''regnet'''  )
           # classification head
           UpperCamelCase__    :   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(_lowerCAmelCase  )
     @add_code_sample_docstrings(
         checkpoint=_IMAGE_CLASS_CHECKPOINT						,						output_type=_lowerCAmelCase						,						config_class=_CONFIG_FOR_DOC						,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT						,						)
     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE = None						,						__SCREAMING_SNAKE_CASE=False						,						)  ->   Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:




           """simple docstring"""







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

           UpperCamelCase__    :   Tuple					=				self.regnet(
               _lowerCAmelCase						,						output_hidden_states=_lowerCAmelCase						,						return_dict=_lowerCAmelCase						,						training=_lowerCAmelCase  )

           UpperCamelCase__    :   Union[str, Any]					=				outputs.pooler_output if return_dict else outputs[1]

           UpperCamelCase__    :   Dict					=				self.classifier[0](_lowerCAmelCase  )
           UpperCamelCase__    :   Optional[int]					=				self.classifier[1](_lowerCAmelCase  )

           UpperCamelCase__    :   List[str]					=				None if labels is None else self.hf_compute_loss(labels=_lowerCAmelCase						,						logits=_lowerCAmelCase  )

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

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




import argparse
import json
from typing import List

from ltp import LTP

from transformers.models.bert.tokenization_bert import BertTokenizer


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      # This defines a "chinese character" as anything in the CJK Unicode block:
      #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
      #
      # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
      # despite its name. The modern Korean Hangul alphabet is a different block,
      # as is Japanese Hiragana and Katakana. Those alphabets are used to write
      # space-separated words, so they are not treated specially and handled
      # like the all of the other languages.
      if (
          (cp >= 0x4e_00 and cp <= 0x9f_ff)
          or (cp >= 0x34_00 and cp <= 0x4d_bf)  #
          or (cp >= 0x2_00_00 and cp <= 0x2_a6_df)  #
          or (cp >= 0x2_a7_00 and cp <= 0x2_b7_3f)  #
          or (cp >= 0x2_b7_40 and cp <= 0x2_b8_1f)  #
          or (cp >= 0x2_b8_20 and cp <= 0x2_ce_af)  #
          or (cp >= 0xf9_00 and cp <= 0xfa_ff)
          or (cp >= 0x2_f8_00 and cp <= 0x2_fa_1f)  #
      ):  #
            return True

      return False


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      # word like '180' or '身高' or '神'
      for char in word:
            UpperCamelCase__    :   Optional[Any]					=				ord(UpperCamelCase__		)
            if not _is_chinese_char(UpperCamelCase__		):
                  return 0
      return 1


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      UpperCamelCase__    :   Dict					=				set()

      for token in tokens:
            UpperCamelCase__    :   Optional[int]					=				len(UpperCamelCase__		) > 1 and is_chinese(UpperCamelCase__		)
            if chinese_word:
                  word_set.add(UpperCamelCase__		)
      UpperCamelCase__    :   List[Any]					=				list(UpperCamelCase__		)
      return word_list


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__		):
      if not chinese_word_set:
            return bert_tokens
      UpperCamelCase__    :   str					=				max([len(UpperCamelCase__		) for w in chinese_word_set]		)

      UpperCamelCase__    :   Optional[int]					=				bert_tokens
      UpperCamelCase__				,UpperCamelCase__    :   Tuple					=				0, len(UpperCamelCase__		)
      while start < end:
            UpperCamelCase__    :   Any					=				True
            if is_chinese(bert_word[start]		):
                  UpperCamelCase__    :   Any					=				min(end - start ,		UpperCamelCase__		)
                  for i in range(UpperCamelCase__ ,		1 ,		-1		):
                        UpperCamelCase__    :   List[Any]					=				''''''.join(bert_word[start : start + i]		)
                        if whole_word in chinese_word_set:
                              for j in range(start + 1 ,		start + i		):
                                    UpperCamelCase__    :   Optional[Any]					=				'''##''' + bert_word[j]
                              UpperCamelCase__    :   Dict					=				start + i
                              UpperCamelCase__    :   List[str]					=				False
                              break
            if single_word:
                  start += 1
      return bert_word


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__ ,		UpperCamelCase__		):
      UpperCamelCase__    :   str					=				[]

      for i in range(0 ,		len(UpperCamelCase__		) ,		1_0_0		):
            UpperCamelCase__    :   Optional[Any]					=				ltp_tokenizer.pipeline(lines[i : i + 1_0_0] ,		tasks=['''cws''']		).cws
            UpperCamelCase__    :   List[str]					=				[get_chinese_word(UpperCamelCase__		) for r in res]
            ltp_res.extend(UpperCamelCase__		)
      assert len(UpperCamelCase__		) == len(UpperCamelCase__		)

      UpperCamelCase__    :   str					=				[]
      for i in range(0 ,		len(UpperCamelCase__		) ,		1_0_0		):
            UpperCamelCase__    :   Dict					=				bert_tokenizer(lines[i : i + 1_0_0] ,		add_special_tokens=UpperCamelCase__ ,		truncation=UpperCamelCase__ ,		max_length=5_1_2		)
            bert_res.extend(res['''input_ids''']		)
      assert len(UpperCamelCase__		) == len(UpperCamelCase__		)

      UpperCamelCase__    :   Any					=				[]
      for input_ids, chinese_word in zip(UpperCamelCase__ ,		UpperCamelCase__		):
            UpperCamelCase__    :   Union[str, Any]					=				[]
            for id in input_ids:
                  UpperCamelCase__    :   List[str]					=				bert_tokenizer._convert_id_to_token(UpperCamelCase__		)
                  input_tokens.append(UpperCamelCase__		)
            UpperCamelCase__    :   Optional[Any]					=				add_sub_symbol(UpperCamelCase__ ,		UpperCamelCase__		)
            UpperCamelCase__    :   Any					=				[]
            # We only save pos of chinese subwords start with ##, which mean is part of a whole word.
            for i, token in enumerate(UpperCamelCase__		):
                  if token[:2] == "##":
                        UpperCamelCase__    :   Union[str, Any]					=				token[2:]
                        # save chinese tokens' pos
                        if len(UpperCamelCase__		) == 1 and _is_chinese_char(ord(UpperCamelCase__		)		):
                              ref_id.append(UpperCamelCase__		)
            ref_ids.append(UpperCamelCase__		)

      assert len(UpperCamelCase__		) == len(UpperCamelCase__		)

      return ref_ids


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      # For Chinese (Ro)Bert, the best result is from : RoBERTa-wwm-ext (https://github.com/ymcui/Chinese-BERT-wwm)
      # If we want to fine-tune these model, we have to use same tokenizer : LTP (https://github.com/HIT-SCIR/ltp)
      with open(args.file_name ,		'''r''' ,		encoding='''utf-8'''		) as f:
            UpperCamelCase__    :   List[Any]					=				f.readlines()
      UpperCamelCase__    :   Tuple					=				[line.strip() for line in data if len(UpperCamelCase__		) > 0 and not line.isspace()]  # avoid delimiter like '\u2029'
      UpperCamelCase__    :   int					=				LTP(args.ltp		)  # faster in GPU device
      UpperCamelCase__    :   int					=				BertTokenizer.from_pretrained(args.bert		)

      UpperCamelCase__    :   Dict					=				prepare_ref(UpperCamelCase__ ,		UpperCamelCase__ ,		UpperCamelCase__		)

      with open(args.save_path ,		'''w''' ,		encoding='''utf-8'''		) as f:
            UpperCamelCase__    :   Tuple					=				[json.dumps(UpperCamelCase__		) + '''\n''' for ref in ref_ids]
            f.writelines(UpperCamelCase__		)


if __name__ == "__main__":
    lowerCamelCase      =argparse.ArgumentParser(description="prepare_chinese_ref")
    parser.add_argument(
        "--file_name",
        required=False,
        type=str,
        default="./resources/chinese-demo.txt",
        help="file need process, same as training data in lm",
    )
    parser.add_argument(
        "--ltp",
        required=False,
        type=str,
        default="./resources/ltp",
        help="resources for LTP tokenizer, usually a path",
    )
    parser.add_argument(
        "--bert",
        required=False,
        type=str,
        default="./resources/robert",
        help="resources for Bert tokenizer",
    )
    parser.add_argument(
        "--save_path",
        required=False,
        type=str,
        default="./resources/ref.txt",
        help="path to save res",
    )

    lowerCamelCase      =parser.parse_args()
    main(args)


