




'''simple docstring'''



from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


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

# General docstring
UpperCAmelCase_				:     int     			=						'''ResNetConfig'''

# Base docstring
UpperCAmelCase_				:     Tuple     			=						'''microsoft/resnet-50'''
UpperCAmelCase_				:     Dict     			=						[1, 2_0_4_8, 7, 7]

# Image classification docstring
UpperCAmelCase_				:     Any     			=						'''microsoft/resnet-50'''
UpperCAmelCase_				:     Union[str, Any]     			=						'''tiger cat'''

UpperCAmelCase_				:     Any     			=						[
    '''microsoft/resnet-50''',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]





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


       def __init__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE = 3			,     __SCREAMING_SNAKE_CASE = 1			,     __SCREAMING_SNAKE_CASE = "relu"			,     **__SCREAMING_SNAKE_CASE     )					->     None:



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     kernel_size // 2
            __snake_case			    =     tf.keras.layers.ConvaD(
                __SCREAMING_SNAKE_CASE			,     kernel_size=__SCREAMING_SNAKE_CASE			,     strides=__SCREAMING_SNAKE_CASE			,     padding='''valid'''			,     use_bias=__SCREAMING_SNAKE_CASE			,     name='''convolution'''     )
            # Use same default momentum and epsilon as PyTorch equivalent
            __snake_case			    =     tf.keras.layers.BatchNormalization(epsilon=1E-5			,     momentum=0.9			,     name='''normalization'''     )
            __snake_case			    =     ACTaFN[activation] if activation is not None else tf.keras.layers.Activation('''linear'''     )


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



            '''simple docstring'''




            __snake_case			    =     __snake_case			    =     (self.pad_value, self.pad_value)
            __snake_case			    =     tf.pad(__SCREAMING_SNAKE_CASE			,     [(0, 0), height_pad, width_pad, (0, 0)]     )
            __snake_case			    =     self.conv(__SCREAMING_SNAKE_CASE     )
            return hidden_state




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



            '''simple docstring'''




            __snake_case			    =     self.convolution(__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.normalization(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.activation(__SCREAMING_SNAKE_CASE     )
            return hidden_state





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


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



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     TFResNetConvLayer(
                config.embedding_size			,     kernel_size=7			,     stride=2			,     activation=config.hidden_act			,     name='''embedder'''			,     )
            __snake_case			    =     tf.keras.layers.MaxPoolaD(pool_size=3			,     strides=2			,     padding='''valid'''			,     name='''pooler'''     )
            __snake_case			    =     config.num_channels




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



            '''simple docstring'''




            __snake_case    ,			__snake_case    ,			__snake_case    ,			__snake_case			    =     shape_list(__SCREAMING_SNAKE_CASE     )
            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.'''     )
            __snake_case			    =     pixel_values
            __snake_case			    =     self.embedder(__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     tf.pad(__SCREAMING_SNAKE_CASE			,     [[0, 0], [1, 1], [1, 1], [0, 0]]     )
            __snake_case			    =     self.pooler(__SCREAMING_SNAKE_CASE     )
            return hidden_state





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


       def __init__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE = 2			,     **__SCREAMING_SNAKE_CASE     )					->     None:



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     tf.keras.layers.ConvaD(
                __SCREAMING_SNAKE_CASE			,     kernel_size=1			,     strides=__SCREAMING_SNAKE_CASE			,     use_bias=__SCREAMING_SNAKE_CASE			,     name='''convolution'''     )
            # Use same default momentum and epsilon as PyTorch equivalent
            __snake_case			    =     tf.keras.layers.BatchNormalization(epsilon=1E-5			,     momentum=0.9			,     name='''normalization'''     )




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



            '''simple docstring'''




            __snake_case			    =     x
            __snake_case			    =     self.convolution(__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.normalization(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            return hidden_state





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


       def __init__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE = 1			,     __SCREAMING_SNAKE_CASE = "relu"			,     **__SCREAMING_SNAKE_CASE     )					->     None:



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     in_channels != out_channels or stride != 1
            __snake_case			    =     TFResNetConvLayer(__SCREAMING_SNAKE_CASE			,     stride=__SCREAMING_SNAKE_CASE			,     name='''layer.0'''     )
            __snake_case			    =     TFResNetConvLayer(__SCREAMING_SNAKE_CASE			,     activation=__SCREAMING_SNAKE_CASE			,     name='''layer.1'''     )
            __snake_case			    =     (
                TFResNetShortCut(__SCREAMING_SNAKE_CASE			,     stride=__SCREAMING_SNAKE_CASE			,     name='''shortcut'''     )
                if should_apply_shortcut
                else tf.keras.layers.Activation('''linear'''			,     name='''shortcut'''     )
            )
            __snake_case			    =     ACTaFN[activation]




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



            '''simple docstring'''




            __snake_case			    =     hidden_state
            __snake_case			    =     self.conva(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.conva(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.shortcut(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            hidden_state += residual
            __snake_case			    =     self.activation(__SCREAMING_SNAKE_CASE     )
            return hidden_state





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


       def __init__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE = 1			,     __SCREAMING_SNAKE_CASE = "relu"			,     __SCREAMING_SNAKE_CASE = 4			,     **__SCREAMING_SNAKE_CASE			,     )					->     None:



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     in_channels != out_channels or stride != 1
            __snake_case			    =     out_channels // reduction
            __snake_case			    =     TFResNetConvLayer(__SCREAMING_SNAKE_CASE			,     kernel_size=1			,     name='''layer.0'''     )
            __snake_case			    =     TFResNetConvLayer(__SCREAMING_SNAKE_CASE			,     stride=__SCREAMING_SNAKE_CASE			,     name='''layer.1'''     )
            __snake_case			    =     TFResNetConvLayer(__SCREAMING_SNAKE_CASE			,     kernel_size=1			,     activation=__SCREAMING_SNAKE_CASE			,     name='''layer.2'''     )
            __snake_case			    =     (
                TFResNetShortCut(__SCREAMING_SNAKE_CASE			,     stride=__SCREAMING_SNAKE_CASE			,     name='''shortcut'''     )
                if should_apply_shortcut
                else tf.keras.layers.Activation('''linear'''			,     name='''shortcut'''     )
            )
            __snake_case			    =     ACTaFN[activation]




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



            '''simple docstring'''




            __snake_case			    =     hidden_state
            __snake_case			    =     self.conva(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.conva(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.conva(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.shortcut(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            hidden_state += residual
            __snake_case			    =     self.activation(__SCREAMING_SNAKE_CASE     )
            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     )					->     None:



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     TFResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else TFResNetBasicLayer

            __snake_case			    =     [layer(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     stride=__SCREAMING_SNAKE_CASE			,     activation=config.hidden_act			,     name='''layers.0'''     )]
            layers += [
                layer(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     activation=config.hidden_act			,     name=F'''layers.{i + 1}'''     )
                for i in range(depth - 1     )
            ]
            __snake_case			    =     layers




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



            '''simple docstring'''




            for layer in self.stage_layers:
                 __snake_case			    =     layer(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )
            return hidden_state





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


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



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )
            # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
            __snake_case			    =     [
                TFResNetStage(
                    __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'''			,     )
            ]
            for i, (in_channels, out_channels, depth) in enumerate(
                zip(config.hidden_sizes			,     config.hidden_sizes[1:]			,     config.depths[1:]     )     ):
                 self.stages.append(TFResNetStage(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     depth=__SCREAMING_SNAKE_CASE			,     name=F'''stages.{i + 1}'''     )     )




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



            '''simple docstring'''




            __snake_case			    =     () if output_hidden_states else None

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

                 __snake_case			    =     stage_module(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )

            if output_hidden_states:
                 __snake_case			    =     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     )





class 			lowerCAmelCase      (       __lowerCAmelCase):
       __lowercase     :					int       				=	ResNetConfig
       __lowercase     :					Optional[Any]       				=	'''resnet'''
       __lowercase     :					Any       				=	'''pixel_values'''


       @property
       def 					lowerCAmelCase					(      self     )					->     Tuple:



            '''simple docstring'''




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





UpperCAmelCase_				:     Any     			=						R'''
    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.

    Parameters:
        config ([`ResNetConfig`]): 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.
'''


UpperCAmelCase_				:     Optional[int]     			=						R'''
    Args:
        pixel_values (`tf.Tensor` 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 [`~utils.ModelOutput`] instead of a plain tuple.
'''





@keras_serializable
class 			lowerCAmelCase      (       tf.keras.layers.Layer):
       __lowercase     :					Tuple       				=	ResNetConfig


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



            '''simple docstring'''




            super().__init__(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     config
            __snake_case			    =     TFResNetEmbeddings(__SCREAMING_SNAKE_CASE			,     name='''embedder'''     )
            __snake_case			    =     TFResNetEncoder(__SCREAMING_SNAKE_CASE			,     name='''encoder'''     )
            __snake_case			    =     tf.keras.layers.GlobalAveragePoolingaD(keepdims=__SCREAMING_SNAKE_CASE     )




       @unpack_inputs
       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = False			,     )					->     Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:



            '''simple docstring'''




            __snake_case			    =     (
                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
            )
            __snake_case			    =     return_dict if return_dict is not None else self.config.use_return_dict

            # TF 2.0 image layers can't use NCHW format when running on CPU.
            # We transpose to NHWC format and then transpose back after the full forward pass.
            # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
            __snake_case			    =     tf.transpose(__SCREAMING_SNAKE_CASE			,     perm=[0, 2, 3, 1]     )
            __snake_case			    =     self.embedder(__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     self.encoder(
                __SCREAMING_SNAKE_CASE			,     output_hidden_states=__SCREAMING_SNAKE_CASE			,     return_dict=__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     encoder_outputs[0]

            __snake_case			    =     self.pooler(__SCREAMING_SNAKE_CASE     )

            # Transpose all the outputs to the NCHW format
            # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
            __snake_case			    =     tf.transpose(__SCREAMING_SNAKE_CASE			,     (0, 3, 1, 2)     )
            __snake_case			    =     tf.transpose(__SCREAMING_SNAKE_CASE			,     (0, 3, 1, 2)     )
            __snake_case			    =     ()
            for hidden_state in encoder_outputs[1:]:
                 __snake_case			    =     hidden_states + tuple(tf.transpose(__SCREAMING_SNAKE_CASE			,     (0, 3, 1, 2)     ) for h in hidden_state     )

            if not return_dict:
                 return (last_hidden_state, pooled_output) + hidden_states

            __snake_case			    =     hidden_states if output_hidden_states else None

            return TFBaseModelOutputWithPoolingAndNoAttention(
                last_hidden_state=__SCREAMING_SNAKE_CASE			,     pooler_output=__SCREAMING_SNAKE_CASE			,     hidden_states=__SCREAMING_SNAKE_CASE			,     )





@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.''' ,      __lowerCAmelCase ,      )
class 			lowerCAmelCase      (       __lowerCAmelCase):


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



            '''simple docstring'''




            super().__init__(__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     TFResNetMainLayer(config=__SCREAMING_SNAKE_CASE			,     name='''resnet'''     )




       @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			,     )
       @unpack_inputs
       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = False			,     )					->     Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:



            '''simple docstring'''




            __snake_case			    =     (
                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
            )
            __snake_case			    =     return_dict if return_dict is not None else self.config.use_return_dict

            __snake_case			    =     self.resnet(
                pixel_values=__SCREAMING_SNAKE_CASE			,     output_hidden_states=__SCREAMING_SNAKE_CASE			,     return_dict=__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE			,     )
            return resnet_outputs





@add_start_docstrings(
    '''
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    ''' ,      __lowerCAmelCase ,      )
class 			lowerCAmelCase      (       __lowerCAmelCase ,      __lowerCAmelCase):


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



            '''simple docstring'''




            super().__init__(__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     config.num_labels
            __snake_case			    =     TFResNetMainLayer(__SCREAMING_SNAKE_CASE			,     name='''resnet'''     )
            # classification head
            __snake_case			    =     (
                tf.keras.layers.Dense(config.num_labels			,     name='''classifier.1'''     )
                if config.num_labels > 0
                else tf.keras.layers.Activation('''linear'''			,     name='''classifier.1'''     )
            )


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



            '''simple docstring'''




            __snake_case			    =     tf.keras.layers.Flatten()(__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     self.classifier_layer(__SCREAMING_SNAKE_CASE     )
            return logits




       @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			,     )
       @unpack_inputs
       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = None			,     __SCREAMING_SNAKE_CASE = False			,     )					->     Union[Tuple[tf.Tensor], TFImageClassifierOutputWithNoAttention]:



            '''simple docstring'''




            __snake_case			    =     return_dict if return_dict is not None else self.config.use_return_dict

            __snake_case			    =     self.resnet(
                __SCREAMING_SNAKE_CASE			,     output_hidden_states=__SCREAMING_SNAKE_CASE			,     return_dict=__SCREAMING_SNAKE_CASE			,     training=__SCREAMING_SNAKE_CASE     )

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

            __snake_case			    =     self.classifier(__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     None if labels is None else self.hf_compute_loss(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     )

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

            return TFImageClassifierOutputWithNoAttention(loss=__SCREAMING_SNAKE_CASE			,     logits=__SCREAMING_SNAKE_CASE			,     hidden_states=outputs.hidden_states     )












'''simple docstring'''



import os
from typing import Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ... import AutoTokenizer
from ...utils import logging


UpperCAmelCase_				:     Optional[int]     			=						'''block_records.npy'''


UpperCAmelCase_				:     Any     			=						logging.get_logger(__name__)







def 	_UpperCamelCase       (_lowerCamelCase							:   str      ,		_lowerCamelCase							:   int		)->      np.ndarray:




     '''simple docstring'''





     import tensorflow.compat.va as tf

     __snake_case			    =     tf.data.TFRecordDataset(_lowerCamelCase      ,		buffer_size=5_12 * 10_24 * 10_24		)
     __snake_case			    =     blocks_dataset.batch(_lowerCamelCase      ,		drop_remainder=_lowerCamelCase		)
     __snake_case			    =     next(blocks_dataset.take(1		).as_numpy_iterator()		)

     return np_record





class 			lowerCAmelCase      :


       def __init__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE=2			,     __SCREAMING_SNAKE_CASE=1000			,     __SCREAMING_SNAKE_CASE=100			,     __SCREAMING_SNAKE_CASE=10_0000			,     )					->     Optional[int]:



            '''simple docstring'''




            from scann.scann_ops.py.scann_ops_pybind import builder as Builder

            __snake_case			    =     Builder(db=__SCREAMING_SNAKE_CASE			,     num_neighbors=__SCREAMING_SNAKE_CASE			,     distance_measure='''dot_product'''     )
            __snake_case			    =     builder.tree(
                num_leaves=__SCREAMING_SNAKE_CASE			,     num_leaves_to_search=__SCREAMING_SNAKE_CASE			,     training_sample_size=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     builder.score_ah(dimensions_per_block=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     builder.build()




       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE     )					->     Any:



            '''simple docstring'''




            __snake_case    ,			__snake_case			    =     self.searcher.search_batched(question_projection.detach().cpu()     )
            return retrieved_block_ids.astype('''int64'''     )





class 			lowerCAmelCase      :


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



            '''simple docstring'''




            super().__init__()
            __snake_case			    =     block_records
            __snake_case			    =     tokenizer


       def __call__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE=None			,     __SCREAMING_SNAKE_CASE="pt"     )					->     Union[str, Any]:



            '''simple docstring'''




            __snake_case			    =     np.take(self.block_records			,     indices=__SCREAMING_SNAKE_CASE			,     axis=0     )

            __snake_case			    =     self.tokenizer.decode(question_input_ids[0]			,     skip_special_tokens=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     []
            __snake_case			    =     []
            for retrieved_block in retrieved_blocks:
                 text.append(__SCREAMING_SNAKE_CASE     )
                 text_pair.append(retrieved_block.decode()     )

            __snake_case			    =     self.tokenizer(
                __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     padding=__SCREAMING_SNAKE_CASE			,     truncation=__SCREAMING_SNAKE_CASE			,     return_special_tokens_mask=__SCREAMING_SNAKE_CASE			,     max_length=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     concat_inputs.convert_to_tensors(__SCREAMING_SNAKE_CASE     )

            if answer_ids is not None:
                 return self.block_has_answer(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     ) + (concat_inputs_tensors,)
            else:
                 return (None, None, None, concat_inputs_tensors)


       @classmethod
       def 					lowerCAmelCase					(      cls			,     __SCREAMING_SNAKE_CASE			,     *__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )					->     Optional[int]:



            '''simple docstring'''




            if os.path.isdir(__SCREAMING_SNAKE_CASE     ):
                 __snake_case			    =     os.path.join(__SCREAMING_SNAKE_CASE			,     _REALM_BLOCK_RECORDS_FILENAME     )
            else:
                 __snake_case			    =     hf_hub_download(
                     repo_id=__SCREAMING_SNAKE_CASE			,     filename=_REALM_BLOCK_RECORDS_FILENAME			,     **__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     np.load(__SCREAMING_SNAKE_CASE			,     allow_pickle=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE			,     *__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )

            return cls(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     )


       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE     )					->     Any:



            '''simple docstring'''




            np.save(os.path.join(__SCREAMING_SNAKE_CASE			,     _REALM_BLOCK_RECORDS_FILENAME     )			,     self.block_records     )
            # save tokenizer
            self.tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE     )




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



            '''simple docstring'''




            __snake_case			    =     []
            __snake_case			    =     []
            __snake_case			    =     []
            __snake_case			    =     0

            for input_id in concat_inputs.input_ids:
                 __snake_case			    =     input_id.tolist()
                 # Check answers between two [SEP] tokens
                 __snake_case			    =     input_id_list.index(self.tokenizer.sep_token_id     )
                 __snake_case			    =     first_sep_idx + 1 + input_id_list[first_sep_idx + 1 :].index(self.tokenizer.sep_token_id     )

                 start_pos.append([]     )
                 end_pos.append([]     )
                 for answer in answer_ids:
                      for idx in range(first_sep_idx + 1			,     __SCREAMING_SNAKE_CASE     ):
                           if answer[0] == input_id_list[idx]:
                                if input_id_list[idx : idx + len(__SCREAMING_SNAKE_CASE     )] == answer:
                                     start_pos[-1].append(__SCREAMING_SNAKE_CASE     )
                                     end_pos[-1].append(idx + len(__SCREAMING_SNAKE_CASE     ) - 1     )

                 if len(start_pos[-1]     ) == 0:
                      has_answers.append(__SCREAMING_SNAKE_CASE     )
                 else:
                      has_answers.append(__SCREAMING_SNAKE_CASE     )
                      if len(start_pos[-1]     ) > max_answers:
                           __snake_case			    =     len(start_pos[-1]     )

        # Pad -1 to max_answers
            for start_pos_, end_pos_ in zip(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     ):
                 if len(start_pos_     ) < max_answers:
                      __snake_case			    =     [-1] * (max_answers - len(start_pos_     ))
                      start_pos_ += padded
                      end_pos_ += padded
            return has_answers, start_pos, end_pos






