
from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


lowercase        =  [
    '''openmmlab/upernet-convnext-tiny''',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
lowercase        =  '''UperNetConfig'''





class       __A(				nn.Module       ):

					def __init__(	self					:     Any   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     Union[int, Tuple[int, int]]   ,					__UpperCamelCase					:     Union[int, Tuple[int, int], str] = 0   ,					__UpperCamelCase					:     bool = False   ,					__UpperCamelCase					:     Union[int, Tuple[int, int]] = 1   ,					):
												super().__init__()
												lowerCamelCase_							  =					nn.Convad(
												    in_channels=__UpperCamelCase   ,					out_channels=__UpperCamelCase   ,					kernel_size=__UpperCamelCase   ,					padding=__UpperCamelCase   ,					bias=__UpperCamelCase   ,					dilation=__UpperCamelCase   ,					)
												lowerCamelCase_							  =					nn.BatchNormad(__UpperCamelCase   )
												lowerCamelCase_							  =					nn.ReLU()



					def    lowercase__      (	self					:     Any   ,					__UpperCamelCase					:     torch.Tensor   ):
												lowerCamelCase_							  =					self.conv(__UpperCamelCase   )
												lowerCamelCase_							  =					self.batch_norm(__UpperCamelCase   )
												lowerCamelCase_							  =					self.activation(__UpperCamelCase   )

												return output





class       __A(				nn.Module       ):

					def __init__(	self					:     Optional[int]   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     int   ):
												super().__init__()
												lowerCamelCase_							  =					[
												    nn.AdaptiveAvgPoolad(__UpperCamelCase   ),
												    UperNetConvModule(__UpperCamelCase   ,					__UpperCamelCase   ,					kernel_size=1   ),
												]
												for i, layer in enumerate(self.layers   ):
																			self.add_module(str(__UpperCamelCase   )   ,					__UpperCamelCase   )



					def    lowercase__      (	self					:     List[Any]   ,					__UpperCamelCase					:     torch.Tensor   ):
												lowerCamelCase_							  =					input
												for layer in self.layers:
																			lowerCamelCase_							  =					layer(__UpperCamelCase   )
												return hidden_state





class       __A(				nn.Module       ):

					def __init__(	self					:     List[str]   ,					__UpperCamelCase					:     Tuple[int, ...]   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     bool   ):
												super().__init__()
												lowerCamelCase_							  =					pool_scales
												lowerCamelCase_							  =					align_corners
												lowerCamelCase_							  =					in_channels
												lowerCamelCase_							  =					channels
												lowerCamelCase_							  =					[]
												for i, pool_scale in enumerate(__UpperCamelCase   ):
																			lowerCamelCase_							  =					UperNetPyramidPoolingBlock(pool_scale=__UpperCamelCase   ,					in_channels=__UpperCamelCase   ,					channels=__UpperCamelCase   )
																			self.blocks.append(__UpperCamelCase   )
																			self.add_module(str(__UpperCamelCase   )   ,					__UpperCamelCase   )



					def    lowercase__      (	self					:     str   ,					__UpperCamelCase					:     torch.Tensor   ):
												lowerCamelCase_							  =					[]
												for ppm in self.blocks:
																			lowerCamelCase_							  =					ppm(__UpperCamelCase   )
																			lowerCamelCase_							  =					nn.functional.interpolate(
																			    __UpperCamelCase   ,					size=x.size()[2:]   ,					mode="""bilinear"""   ,					align_corners=self.align_corners   )
																			ppm_outs.append(__UpperCamelCase   )
												return ppm_outs





class       __A(				nn.Module       ):

					def __init__(	self					:     str   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     List[Any]   ):
												super().__init__()

												lowerCamelCase_							  =					config
												lowerCamelCase_							  =					config.pool_scales  # e.g. (1, 2, 3, 6)
												lowerCamelCase_							  =					in_channels
												lowerCamelCase_							  =					config.hidden_size
												lowerCamelCase_							  =					False
												lowerCamelCase_							  =					nn.Convad(self.channels   ,					config.num_labels   ,					kernel_size=1   )

												# PSP Module
												lowerCamelCase_							  =					UperNetPyramidPoolingModule(
												    self.pool_scales   ,					self.in_channels[-1]   ,					self.channels   ,					align_corners=self.align_corners   ,					)
												lowerCamelCase_							  =					UperNetConvModule(
												    self.in_channels[-1] + len(self.pool_scales   ) * self.channels   ,					self.channels   ,					kernel_size=3   ,					padding=1   ,					)
												# FPN Module
												lowerCamelCase_							  =					nn.ModuleList()
												lowerCamelCase_							  =					nn.ModuleList()
												for in_channels in self.in_channels[:-1]:  # skip the top layer
																			lowerCamelCase_							  =					UperNetConvModule(__UpperCamelCase   ,					self.channels   ,					kernel_size=1   )
																			lowerCamelCase_							  =					UperNetConvModule(self.channels   ,					self.channels   ,					kernel_size=3   ,					padding=1   )
																			self.lateral_convs.append(__UpperCamelCase   )
																			self.fpn_convs.append(__UpperCamelCase   )

												lowerCamelCase_							  =					UperNetConvModule(
												    len(self.in_channels   ) * self.channels   ,					self.channels   ,					kernel_size=3   ,					padding=1   ,					)

					def    lowercase__      (	self					:     Optional[Any]   ):
												self.apply(self._init_weights   )

					def    lowercase__      (	self					:     Optional[Any]   ,					__UpperCamelCase					:     Optional[int]   ):
												if isinstance(__UpperCamelCase   ,					nn.Convad   ):
																			module.weight.data.normal_(mean=0.0   ,					std=self.config.initializer_range   )
																			if module.bias is not None:
																										module.bias.data.zero_()

					def    lowercase__      (	self					:     Optional[int]   ,					__UpperCamelCase					:     Union[str, Any]   ):
												lowerCamelCase_							  =					inputs[-1]
												lowerCamelCase_							  =					[x]
												psp_outs.extend(self.psp_modules(__UpperCamelCase   )   )
												lowerCamelCase_							  =					torch.cat(__UpperCamelCase   ,					dim=1   )
												lowerCamelCase_							  =					self.bottleneck(__UpperCamelCase   )

												return output



					def    lowercase__      (	self					:     Tuple   ,					__UpperCamelCase					:     torch.Tensor   ):
												# build laterals
												lowerCamelCase_							  =					[lateral_conv(encoder_hidden_states[i]   ) for i, lateral_conv in enumerate(self.lateral_convs   )]

												laterals.append(self.psp_forward(__UpperCamelCase   )   )

												# build top-down path
												lowerCamelCase_							  =					len(__UpperCamelCase   )
												for i in range(used_backbone_levels - 1   ,					0   ,					-1   ):
																			lowerCamelCase_							  =					laterals[i - 1].shape[2:]
																			lowerCamelCase_							  =					laterals[i - 1] + nn.functional.interpolate(
																			    laterals[i]   ,					size=__UpperCamelCase   ,					mode="""bilinear"""   ,					align_corners=self.align_corners   )

												# build outputs
												lowerCamelCase_							  =					[self.fpn_convs[i](laterals[i]   ) for i in range(used_backbone_levels - 1   )]
												# append psp feature
												fpn_outs.append(laterals[-1]   )

												for i in range(used_backbone_levels - 1   ,					0   ,					-1   ):
																			lowerCamelCase_							  =					nn.functional.interpolate(
																			    fpn_outs[i]   ,					size=fpn_outs[0].shape[2:]   ,					mode="""bilinear"""   ,					align_corners=self.align_corners   )
												lowerCamelCase_							  =					torch.cat(__UpperCamelCase   ,					dim=1   )
												lowerCamelCase_							  =					self.fpn_bottleneck(__UpperCamelCase   )
												lowerCamelCase_							  =					self.classifier(__UpperCamelCase   )

												return output





class       __A(				nn.Module       ):

					def __init__(	self					:     Union[str, Any]   ,					__UpperCamelCase					:     Union[str, Any]   ,					__UpperCamelCase					:     int = 2   ,					__UpperCamelCase					:     int = 3   ,					__UpperCamelCase					:     Union[int, Tuple[int, int]] = 1   ):
												super().__init__()

												lowerCamelCase_							  =					config
												lowerCamelCase_							  =					config.auxiliary_in_channels
												lowerCamelCase_							  =					config.auxiliary_channels
												lowerCamelCase_							  =					config.auxiliary_num_convs
												lowerCamelCase_							  =					config.auxiliary_concat_input
												lowerCamelCase_							  =					in_index

												lowerCamelCase_							  =					(kernel_size // 2) * dilation
												lowerCamelCase_							  =					[]
												convs.append(
												    UperNetConvModule(
												        self.in_channels   ,					self.channels   ,					kernel_size=__UpperCamelCase   ,					padding=__UpperCamelCase   ,					dilation=__UpperCamelCase   )   )
												for i in range(self.num_convs - 1   ):
																			convs.append(
																			    UperNetConvModule(
																			        self.channels   ,					self.channels   ,					kernel_size=__UpperCamelCase   ,					padding=__UpperCamelCase   ,					dilation=__UpperCamelCase   )   )
												if self.num_convs == 0:
																			lowerCamelCase_							  =					nn.Identity()
												else:
																			lowerCamelCase_							  =					nn.Sequential(*__UpperCamelCase   )
												if self.concat_input:
																			lowerCamelCase_							  =					UperNetConvModule(
																			    self.in_channels + self.channels   ,					self.channels   ,					kernel_size=__UpperCamelCase   ,					padding=kernel_size // 2   )

												lowerCamelCase_							  =					nn.Convad(self.channels   ,					config.num_labels   ,					kernel_size=1   )

					def    lowercase__      (	self					:     List[Any]   ):
												self.apply(self._init_weights   )

					def    lowercase__      (	self					:     int   ,					__UpperCamelCase					:     Tuple   ):
												if isinstance(__UpperCamelCase   ,					nn.Convad   ):
																			module.weight.data.normal_(mean=0.0   ,					std=self.config.initializer_range   )
																			if module.bias is not None:
																										module.bias.data.zero_()



					def    lowercase__      (	self					:     Optional[Any]   ,					__UpperCamelCase					:     torch.Tensor   ):
												# just take the relevant feature maps
												lowerCamelCase_							  =					encoder_hidden_states[self.in_index]
												lowerCamelCase_							  =					self.convs(__UpperCamelCase   )
												if self.concat_input:
																			lowerCamelCase_							  =					self.conv_cat(torch.cat([hidden_states, output]   ,					dim=1   )   )
												lowerCamelCase_							  =					self.classifier(__UpperCamelCase   )
												return output





class       __A(				UpperCAmelCase       ):
					SCREAMING_SNAKE_CASE			    =    UperNetConfig
					SCREAMING_SNAKE_CASE			    =    '''pixel_values'''
					SCREAMING_SNAKE_CASE			    =    True

					def    lowercase__      (	self					:     str   ,					__UpperCamelCase					:     Dict   ):
												if isinstance(__UpperCamelCase   ,					__UpperCamelCase   ):
																			module.backbone.init_weights()
																			module.decode_head.init_weights()
																			module.auxiliary_head.init_weights()

					def    lowercase__      (	self					:     Optional[int]   ):
												self.backbone.init_weights()
												self.decode_head.init_weights()
												self.auxiliary_head.init_weights()



					def    lowercase__      (	self					:     str   ,					__UpperCamelCase					:     Tuple   ,					__UpperCamelCase					:     Optional[Any]=False   ):
												if isinstance(__UpperCamelCase   ,					__UpperCamelCase   ):
																			lowerCamelCase_							  =					value
lowercase        =  r'''
    Parameters:
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.
        config ([`UperNetConfig`]): 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        =  r'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See
            `attentions` under returned tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers of the backbone. 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(
    '''UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes.'''     ,		UpperCAmelCase     ,		)
class       __A(				UpperCAmelCase       ):

					def __init__(	self					:     str   ,					__UpperCamelCase					:     Any   ):
												super().__init__(__UpperCamelCase   )

												lowerCamelCase_							  =					AutoBackbone.from_config(config.backbone_config   )

												# Semantic segmentation head(s)
												lowerCamelCase_							  =					UperNetHead(__UpperCamelCase   ,					in_channels=self.backbone.channels   )
												lowerCamelCase_							  =					UperNetFCNHead(__UpperCamelCase   ) if config.use_auxiliary_head else None

												# Initialize weights and apply final processing
												self.post_init()



					@add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format("""batch_size, sequence_length"""   )   )
					@replace_return_docstrings(output_type=__UpperCamelCase   ,					config_class=_CONFIG_FOR_DOC   )
					def    lowercase__      (	self					:     int   ,					__UpperCamelCase					:     Optional[torch.Tensor] = None   ,					__UpperCamelCase					:     Optional[bool] = None   ,					__UpperCamelCase					:     Optional[bool] = None   ,					__UpperCamelCase					:     Optional[torch.Tensor] = None   ,					__UpperCamelCase					:     Optional[bool] = None   ,					):
												lowerCamelCase_							  =					return_dict if return_dict is not None else self.config.use_return_dict
												lowerCamelCase_							  =					(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												lowerCamelCase_							  =					output_attentions if output_attentions is not None else self.config.output_attentions

												lowerCamelCase_							  =					self.backbone.forward_with_filtered_kwargs(
												    __UpperCamelCase   ,					output_hidden_states=__UpperCamelCase   ,					output_attentions=__UpperCamelCase   )
												lowerCamelCase_							  =					outputs.feature_maps

												lowerCamelCase_							  =					self.decode_head(__UpperCamelCase   )
												lowerCamelCase_							  =					nn.functional.interpolate(__UpperCamelCase   ,					size=pixel_values.shape[2:]   ,					mode="""bilinear"""   ,					align_corners=__UpperCamelCase   )

												lowerCamelCase_							  =					None
												if self.auxiliary_head is not None:
																			lowerCamelCase_							  =					self.auxiliary_head(__UpperCamelCase   )
																			lowerCamelCase_							  =					nn.functional.interpolate(
																			    __UpperCamelCase   ,					size=pixel_values.shape[2:]   ,					mode="""bilinear"""   ,					align_corners=__UpperCamelCase   )

												lowerCamelCase_							  =					None
												if labels is not None:
																			if self.config.num_labels == 1:
																										raise ValueError("""The number of labels should be greater than one"""   )
																			else:
																										# compute weighted loss
																										lowerCamelCase_							  =					CrossEntropyLoss(ignore_index=self.config.loss_ignore_index   )
																										lowerCamelCase_							  =					loss_fct(__UpperCamelCase   ,					__UpperCamelCase   )
																										lowerCamelCase_							  =					loss_fct(__UpperCamelCase   ,					__UpperCamelCase   )
																										lowerCamelCase_							  =					main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

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

												return SemanticSegmenterOutput(
												    loss=__UpperCamelCase   ,					logits=__UpperCamelCase   ,					hidden_states=outputs.hidden_states   ,					attentions=outputs.attentions   ,					)







from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import numpy
		import tensorflow as tf

		from transformers import (
		    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    BertConfig,
		    DPRConfig,
		    TFDPRContextEncoder,
		    TFDPRQuestionEncoder,
		    TFDPRReader,
		)





class       __A:

					def __init__(	self					:     str   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     Any=1_3   ,					__UpperCamelCase					:     Dict=7   ,					__UpperCamelCase					:     int=True   ,					__UpperCamelCase					:     Optional[Any]=True   ,					__UpperCamelCase					:     Optional[Any]=True   ,					__UpperCamelCase					:     Dict=True   ,					__UpperCamelCase					:     Dict=9_9   ,					__UpperCamelCase					:     Optional[int]=3_2   ,					__UpperCamelCase					:     List[Any]=2   ,					__UpperCamelCase					:     Optional[int]=4   ,					__UpperCamelCase					:     List[Any]=3_7   ,					__UpperCamelCase					:     Optional[Any]="gelu"   ,					__UpperCamelCase					:     List[str]=0.1   ,					__UpperCamelCase					:     Any=0.1   ,					__UpperCamelCase					:     int=5_1_2   ,					__UpperCamelCase					:     List[Any]=1_6   ,					__UpperCamelCase					:     List[Any]=2   ,					__UpperCamelCase					:     List[str]=0.02   ,					__UpperCamelCase					:     Any=3   ,					__UpperCamelCase					:     int=4   ,					__UpperCamelCase					:     Optional[Any]=None   ,					__UpperCamelCase					:     Union[str, Any]=0   ,					):
												lowerCamelCase_							  =					parent
												lowerCamelCase_							  =					batch_size
												lowerCamelCase_							  =					seq_length
												lowerCamelCase_							  =					is_training
												lowerCamelCase_							  =					use_input_mask
												lowerCamelCase_							  =					use_token_type_ids
												lowerCamelCase_							  =					use_labels
												lowerCamelCase_							  =					vocab_size
												lowerCamelCase_							  =					hidden_size
												lowerCamelCase_							  =					num_hidden_layers
												lowerCamelCase_							  =					num_attention_heads
												lowerCamelCase_							  =					intermediate_size
												lowerCamelCase_							  =					hidden_act
												lowerCamelCase_							  =					hidden_dropout_prob
												lowerCamelCase_							  =					attention_probs_dropout_prob
												lowerCamelCase_							  =					max_position_embeddings
												lowerCamelCase_							  =					type_vocab_size
												lowerCamelCase_							  =					type_sequence_label_size
												lowerCamelCase_							  =					initializer_range
												lowerCamelCase_							  =					num_labels
												lowerCamelCase_							  =					num_choices
												lowerCamelCase_							  =					scope
												lowerCamelCase_							  =					projection_dim

					def    lowercase__      (	self					:     Union[str, Any]   ):
												lowerCamelCase_							  =					ids_tensor([self.batch_size, self.seq_length]   ,					self.vocab_size   )

												lowerCamelCase_							  =					None
												if self.use_input_mask:
																			# follow test_modeling_tf_ctrl.py
																			lowerCamelCase_							  =					random_attention_mask([self.batch_size, self.seq_length]   )

												lowerCamelCase_							  =					None
												if self.use_token_type_ids:
																			lowerCamelCase_							  =					ids_tensor([self.batch_size, self.seq_length]   ,					self.type_vocab_size   )

												lowerCamelCase_							  =					None
												lowerCamelCase_							  =					None
												lowerCamelCase_							  =					None
												if self.use_labels:
																			lowerCamelCase_							  =					ids_tensor([self.batch_size]   ,					self.type_sequence_label_size   )
																			lowerCamelCase_							  =					ids_tensor([self.batch_size, self.seq_length]   ,					self.num_labels   )
																			lowerCamelCase_							  =					ids_tensor([self.batch_size]   ,					self.num_choices   )

												lowerCamelCase_							  =					BertConfig(
												    vocab_size=self.vocab_size   ,					hidden_size=self.hidden_size   ,					num_hidden_layers=self.num_hidden_layers   ,					num_attention_heads=self.num_attention_heads   ,					intermediate_size=self.intermediate_size   ,					hidden_act=self.hidden_act   ,					hidden_dropout_prob=self.hidden_dropout_prob   ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,					max_position_embeddings=self.max_position_embeddings   ,					type_vocab_size=self.type_vocab_size   ,					is_decoder=__UpperCamelCase   ,					initializer_range=self.initializer_range   ,					)
												lowerCamelCase_							  =					DPRConfig(projection_dim=self.projection_dim   ,					**config.to_dict()   )

												return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

					def    lowercase__      (	self					:     Optional[Any]   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     Any   ,					__UpperCamelCase					:     Optional[int]   ,					__UpperCamelCase					:     Optional[int]   ,					__UpperCamelCase					:     List[str]   ,					__UpperCamelCase					:     Union[str, Any]   ):
												lowerCamelCase_							  =					TFDPRContextEncoder(config=__UpperCamelCase   )
												lowerCamelCase_							  =					model(__UpperCamelCase   ,					attention_mask=__UpperCamelCase   ,					token_type_ids=__UpperCamelCase   )
												lowerCamelCase_							  =					model(__UpperCamelCase   ,					token_type_ids=__UpperCamelCase   )
												lowerCamelCase_							  =					model(__UpperCamelCase   )
												self.parent.assertEqual(result.pooler_output.shape   ,					(self.batch_size, self.projection_dim or self.hidden_size)   )

					def    lowercase__      (	self					:     str   ,					__UpperCamelCase					:     List[Any]   ,					__UpperCamelCase					:     Tuple   ,					__UpperCamelCase					:     Optional[int]   ,					__UpperCamelCase					:     Any   ,					__UpperCamelCase					:     Union[str, Any]   ,					__UpperCamelCase					:     List[Any]   ,					__UpperCamelCase					:     Union[str, Any]   ):
												lowerCamelCase_							  =					TFDPRQuestionEncoder(config=__UpperCamelCase   )
												lowerCamelCase_							  =					model(__UpperCamelCase   ,					attention_mask=__UpperCamelCase   ,					token_type_ids=__UpperCamelCase   )
												lowerCamelCase_							  =					model(__UpperCamelCase   ,					token_type_ids=__UpperCamelCase   )
												lowerCamelCase_							  =					model(__UpperCamelCase   )
												self.parent.assertEqual(result.pooler_output.shape   ,					(self.batch_size, self.projection_dim or self.hidden_size)   )

					def    lowercase__      (	self					:     str   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     Tuple   ,					__UpperCamelCase					:     Optional[int]   ,					__UpperCamelCase					:     Tuple   ,					__UpperCamelCase					:     Any   ,					__UpperCamelCase					:     Any   ,					__UpperCamelCase					:     List[str]   ):
												lowerCamelCase_							  =					TFDPRReader(config=__UpperCamelCase   )
												lowerCamelCase_							  =					model(__UpperCamelCase   ,					attention_mask=__UpperCamelCase   )

												self.parent.assertEqual(result.start_logits.shape   ,					(self.batch_size, self.seq_length)   )
												self.parent.assertEqual(result.end_logits.shape   ,					(self.batch_size, self.seq_length)   )
												self.parent.assertEqual(result.relevance_logits.shape   ,					(self.batch_size,)   )



					def    lowercase__      (	self					:     Dict   ):
												lowerCamelCase_							  =					self.prepare_config_and_inputs()
												(
												    (
												    lowerCamelCase_
												)  ,   (
												    lowerCamelCase_
												)  ,   (
												    lowerCamelCase_
												)  ,   (
												    lowerCamelCase_
												)  ,   (
												    lowerCamelCase_
												)  ,   (
												    lowerCamelCase_
												)  ,   (
												    lowerCamelCase_
												)  ,   
												)							  =					config_and_inputs
												lowerCamelCase_							  =					{"""input_ids""": input_ids}
												return config, inputs_dict





@require_tf
class       __A(				UpperCAmelCase     ,		UpperCAmelCase     ,		unittest.TestCase       ):
					SCREAMING_SNAKE_CASE			    =    (
					    (
					        TFDPRContextEncoder,
					        TFDPRQuestionEncoder,
					        TFDPRReader,
					    )
					    if is_tf_available()
					    else ()
					)
					SCREAMING_SNAKE_CASE			    =    {'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

					SCREAMING_SNAKE_CASE			    =    False
					SCREAMING_SNAKE_CASE			    =    False
					SCREAMING_SNAKE_CASE			    =    False
					SCREAMING_SNAKE_CASE			    =    False
					SCREAMING_SNAKE_CASE			    =    False

					def    lowercase__      (	self					:     Dict   ):
												lowerCamelCase_							  =					TFDPRModelTester(self   )
												lowerCamelCase_							  =					ConfigTester(self   ,					config_class=__UpperCamelCase   ,					hidden_size=3_7   )

					def    lowercase__      (	self					:     Optional[Any]   ):
												self.config_tester.run_common_tests()

					def    lowercase__      (	self					:     Any   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_dpr_context_encoder(*__UpperCamelCase   )

					def    lowercase__      (	self					:     Dict   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_dpr_question_encoder(*__UpperCamelCase   )

					def    lowercase__      (	self					:     List[str]   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_dpr_reader(*__UpperCamelCase   )



					@slow
					def    lowercase__      (	self					:     Optional[int]   ):
												for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			lowerCamelCase_							  =					TFDPRContextEncoder.from_pretrained(__UpperCamelCase   )
																			self.assertIsNotNone(__UpperCamelCase   )

												for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			lowerCamelCase_							  =					TFDPRContextEncoder.from_pretrained(__UpperCamelCase   )
																			self.assertIsNotNone(__UpperCamelCase   )

												for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			lowerCamelCase_							  =					TFDPRQuestionEncoder.from_pretrained(__UpperCamelCase   )
																			self.assertIsNotNone(__UpperCamelCase   )

												for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			lowerCamelCase_							  =					TFDPRReader.from_pretrained(__UpperCamelCase   )
																			self.assertIsNotNone(__UpperCamelCase   )





@require_tf
class       __A(				unittest.TestCase       ):

					@slow
					def    lowercase__      (	self					:     Union[str, Any]   ):
												lowerCamelCase_							  =					TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""   )

												lowerCamelCase_							  =					tf.constant(
												    [[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_0_3, 2_0_2_6, 3_8_9_9, 1_0_1_4_0, 1_0_2_9, 1_0_2]]   )  # [CLS] hello, is my dog cute? [SEP]
												lowerCamelCase_							  =					model(__UpperCamelCase   )[0]  # embedding shape = (1, 768)
												# compare the actual values for a slice.
												lowerCamelCase_							  =					tf.constant(
												    [
												        [
												            0.03236253,
												            0.12753335,
												            0.16818509,
												            0.00279786,
												            0.3896933,
												            0.24264945,
												            0.2178971,
												            -0.02335227,
												            -0.08481959,
												            -0.14324117,
												        ]
												    ]   )
												self.assertTrue(numpy.allclose(output[:, :1_0].numpy()   ,					expected_slice.numpy()   ,					atol=1E-4   )   )





