codingutf8 2018 dpr s the hugging face team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tensorflow dpr model for open domain question answering from future import annotations from dataclasses import dataclass from typing import tuple union import tensorflow as tf from modelingtfoutputs import tfbasemodeloutputwithpooling from modelingtfutils import tfmodelinputtype tfpretrainedmodel getinitializer shapelist unpackinputs from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from bert modelingtfbert import tfbertmainlayer from configurationdpr import dprconfig logger logging getloggername configfordoc dprconfig tfdprcontextencoderpretrainedmodelarchivelist facebookdprctxencodersinglenqbase facebookdprctxencodermultisetbase tfdprquestionencoderpretrainedmodelarchivelist facebookdprquestionencodersinglenqbase facebookdprquestionencodermultisetbase tfdprreaderpretrainedmodelarchivelist facebookdprreadersinglenqbase facebookdprreadermultisetbase outputs dataclass class tfdprcontextencoderoutputmodeloutput r class for outputs of tfdprcontextencoder args pooleroutput tf tensor of shape batchsize embeddingssize the dpr encoder outputs the pooleroutput that corresponds to the context representation last layer hiddenstate of the first token of the sequence classification token further processed by a linear layer this output is to be used to embed contexts for nearest neighbors queries with questions embeddings hiddenstates tupletf tensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletf tensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of tf tensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads class for outputs of tfdprquestionencoder args pooleroutput tf tensor of shape batchsize embeddingssize the dpr encoder outputs the pooleroutput that corresponds to the question representation last layer hiddenstate of the first token of the sequence classification token further processed by a linear layer this output is to be used to embed questions for nearest neighbors queries with context embeddings hiddenstates tupletf tensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletf tensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of tf tensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads class for outputs of tfdprreaderencoder args startlogits tf tensor of shape npassages sequencelength logits of the start index of the span for each passage endlogits tf tensor of shape npassages sequencelength logits of the end index of the span for each passage relevancelogits tf tensor of shape npassages outputs of the qa classifier of the dprreader that corresponds to the scores of each passage to answer the question compared to all the other passages hiddenstates tupletf tensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletorch floattensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of torch floattensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads resolve name conflict with tfbertmainlayer instead of tfbertmodel notations n number of questions in a batch m number of passages per questions l sequence length feed encoder compute logits resize pretrainedmodel an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models actual models tfdprencodersinputsdocstring r args inputids numpy array or tf tensor of shape batchsize sequencelength indices of input sequence tokens in the vocabulary to match pretraining dpr input sequence should be formatted with cls and sep tokens as follows a for sequence pairs for a pair titletext for example tokens cls is this jack son ville sep no it is not sep tokentypeids 0 0 0 0 0 0 0 0 1 1 1 1 1 1 b for single sequences for a question for example tokens cls the dog is hairy sep tokentypeids 0 0 0 0 0 0 0 dpr is a model with absolute position embeddings so it s usually advised to pad the inputs on the right rather than the left indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids attentionmask numpy array or tf tensor of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask tokentypeids numpy array or tf tensor of shape batchsize sequencelength optional segment token indices to indicate first and second portions of the inputs indices are selected in 0 1 0 corresponds to a sentence a token 1 corresponds to a sentence b token what are token type ids glossarytokentypeids inputsembeds numpy array or tf tensor of shape batchsize sequencelength hiddensize optional optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation addstartdocstrings the bare dprcontextencoder transformer outputting pooler outputs as context representations tfdprstartdocstring class tfdprcontextencodertfdprpretrainedcontextencoder def initself config dprconfig args kwargs super initconfig args kwargs self ctxencoder tfdprencoderlayerconfig namectxencoder def getinputembeddingsself try return self ctxencoder bertmodel getinputembeddings except attributeerror self build return self ctxencoder bertmodel getinputembeddings unpackinputs addstartdocstringstomodelforwardtfdprencodersinputsdocstring replacereturndocstringsoutputtypetfdprcontextencoderoutput configclassconfigfordoc def call self inputids tfmodelinputtype none none attentionmask tf tensor none none tokentypeids tf tensor none none inputsembeds tf tensor none none outputattentions bool none none outputhiddenstates bool none none returndict bool none none training bool false tfdprcontextencoderoutput tupletf tensor r return examples python from transformers import tfdprcontextencoder dprcontextencodertokenizer tokenizer dprcontextencodertokenizer frompretrainedfacebookdprctxencodersinglenqbase model tfdprcontextencoder frompretrainedfacebookdprctxencodersinglenqbase frompttrue inputids tokenizerhello is my dog cute returntensorstfinputids embeddings modelinputids pooleroutput if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both inputids and inputsembeds at the same time elif inputids is not none inputshape shapelistinputids elif inputsembeds is not none inputshape shapelistinputsembeds 1 else raise valueerroryou have to specify either inputids or inputsembeds if attentionmask is none attentionmask tf onesinputshape dtypetf dtypes int32 if inputids is none else inputids self config padtokenid if tokentypeids is none tokentypeids tf zerosinputshape dtypetf dtypes int32 outputs self questionencoder inputidsinputids attentionmaskattentionmask tokentypeidstokentypeids inputsembedsinputsembeds outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict trainingtraining if not returndict return outputs1 return tfdprquestionencoderoutput pooleroutputoutputs pooleroutput hiddenstatesoutputs hiddenstates attentionsoutputs attentions addstartdocstrings the bare dprreader transformer outputting span predictions tfdprstartdocstring class tfdprreadertfdprpretrainedreader def initself config dprconfig args kwargs super initconfig args kwargs self spanpredictor tfdprspanpredictorlayerconfig namespanpredictor def getinputembeddingsself try return self spanpredictor encoder bertmodel getinputembeddings except attributeerror self build return self spanpredictor encoder bertmodel getinputembeddings unpackinputs addstartdocstringstomodelforwardtfdprreaderinputsdocstring replacereturndocstringsoutputtypetfdprreaderoutput configclassconfigfordoc def call self inputids tfmodelinputtype none none attentionmask tf tensor none none inputsembeds tf tensor none none outputattentions bool none none outputhiddenstates bool none none returndict bool none none training bool false tfdprreaderoutput tupletf tensor r return examples python from transformers import tfdprreader dprreadertokenizer tokenizer dprreadertokenizer frompretrainedfacebookdprreadersinglenqbase model tfdprreader frompretrainedfacebookdprreadersinglenqbase frompttrue encodedinputs tokenizer questionswhat is love titleshaddaway texts what is love is a song recorded by the artist haddaway returntensorstf outputs modelencodedinputs startlogits outputs startlogits endlogits outputs endlogits relevancelogits outputs relevancelogits coding utf 8 2018 dpr s the hugging face team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tensorflow dpr model for open domain question answering outputs class for outputs of tfdprcontextencoder args pooler_output tf tensor of shape batch_size embeddings_size the dpr encoder outputs the pooler_output that corresponds to the context representation last layer hidden state of the first token of the sequence classification token further processed by a linear layer this output is to be used to embed contexts for nearest neighbors queries with questions embeddings hidden_states tuple tf tensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple tf tensor optional returned when output_attentions true is passed or when config output_attentions true tuple of tf tensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads class for outputs of tfdprquestionencoder args pooler_output tf tensor of shape batch_size embeddings_size the dpr encoder outputs the pooler_output that corresponds to the question representation last layer hidden state of the first token of the sequence classification token further processed by a linear layer this output is to be used to embed questions for nearest neighbors queries with context embeddings hidden_states tuple tf tensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple tf tensor optional returned when output_attentions true is passed or when config output_attentions true tuple of tf tensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads class for outputs of tfdprreaderencoder args start_logits tf tensor of shape n_passages sequence_length logits of the start index of the span for each passage end_logits tf tensor of shape n_passages sequence_length logits of the end index of the span for each passage relevance_logits tf tensor of shape n_passages outputs of the qa classifier of the dprreader that corresponds to the scores of each passage to answer the question compared to all the other passages hidden_states tuple tf tensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads resolve name conflict with tfbertmainlayer instead of tfbertmodel notations n number of questions in a batch m number of passages per questions l sequence length feed encoder compute logits resize pretrainedmodel an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models actual models this model inherits from tfpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a tensorflow tf keras model https www tensorflow org api_docs python tf keras model subclass use it as a regular tf 2 0 keras model and refer to the tf 2 0 documentation for all matter related to general usage and behavior tip tensorflow models and layers in transformers accept two formats as input having all inputs as keyword arguments like pytorch models or having all inputs as a list tuple or dict in the first positional argument the reason the second format is supported is that keras methods prefer this format when passing inputs to models and layers because of this support when using methods like model fit things should just work for you just pass your inputs and labels in any format that model fit supports if however you want to use the second format outside of keras methods like fit and predict such as when creating your own layers or models with the keras functional api there are three possibilities you can use to gather all the input tensors in the first positional argument a single tensor with input_ids only and nothing else model input_ids a list of varying length with one or several input tensors in the order given in the docstring model input_ids attention_mask or model input_ids attention_mask token_type_ids a dictionary with one or several input tensors associated to the input names given in the docstring model input_ids input_ids token_type_ids token_type_ids note that when creating models and layers with subclassing https keras io guides making_new_layers_and_models_via_subclassing then you don t need to worry about any of this as you can just pass inputs like you would to any other python function tip parameters config dprconfig 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 args input_ids numpy array or tf tensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary to match pretraining dpr input sequence should be formatted with cls and sep tokens as follows a for sequence pairs for a pair title text for example tokens cls is this jack son ville sep no it is not sep token_type_ids 0 0 0 0 0 0 0 0 1 1 1 1 1 1 b for single sequences for a question for example tokens cls the dog is hairy sep token_type_ids 0 0 0 0 0 0 0 dpr is a model with absolute position embeddings so it s usually advised to pad the inputs on the right rather than the left indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask numpy array or tf tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask token_type_ids numpy array or tf tensor of shape batch_size sequence_length optional segment token indices to indicate first and second portions of the inputs indices are selected in 0 1 0 corresponds to a sentence a token 1 corresponds to a sentence b token what are token type ids glossary token type ids inputs_embeds numpy array or tf tensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation args input_ids numpy array or tf tensor of shapes n_passages sequence_length indices of input sequence tokens in the vocabulary it has to be a sequence triplet with 1 the question and 2 the passages titles and 3 the passages texts to match pretraining dpr input_ids sequence should be formatted with cls and sep with the format cls question token ids sep titles ids sep texts ids dpr is a model with absolute position embeddings so it s usually advised to pad the inputs on the right rather than the left indices can be obtained using dprreadertokenizer see this class documentation for more details attention_mask numpy array or tf tensor of shape n_passages sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask inputs_embeds numpy array or tf tensor of shape n_passages sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation return examples python from transformers import tfdprcontextencoder dprcontextencodertokenizer tokenizer dprcontextencodertokenizer from_pretrained facebook dpr ctx_encoder single nq base model tfdprcontextencoder from_pretrained facebook dpr ctx_encoder single nq base from_pt true input_ids tokenizer hello is my dog cute return_tensors tf input_ids embeddings model input_ids pooler_output return examples python from transformers import tfdprquestionencoder dprquestionencodertokenizer tokenizer dprquestionencodertokenizer from_pretrained facebook dpr question_encoder single nq base model tfdprquestionencoder from_pretrained facebook dpr question_encoder single nq base from_pt true input_ids tokenizer hello is my dog cute return_tensors tf input_ids embeddings model input_ids pooler_output return examples python from transformers import tfdprreader dprreadertokenizer tokenizer dprreadertokenizer from_pretrained facebook dpr reader single nq base model tfdprreader from_pretrained facebook dpr reader single nq base from_pt true encoded_inputs tokenizer questions what is love titles haddaway texts what is love is a song recorded by the artist haddaway return_tensors tf outputs model encoded_inputs start_logits outputs start_logits end_logits outputs end_logits relevance_logits outputs relevance_logits
from __future__ import annotations

from dataclasses import dataclass
from typing import Tuple, Union

import tensorflow as tf

from ...modeling_tf_outputs import TFBaseModelOutputWithPooling
from ...modeling_tf_utils import TFModelInputType, TFPreTrainedModel, get_initializer, shape_list, unpack_inputs
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ..bert.modeling_tf_bert import TFBertMainLayer
from .configuration_dpr import DPRConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "DPRConfig"

TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/dpr-ctx_encoder-single-nq-base",
    "facebook/dpr-ctx_encoder-multiset-base",
]
TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/dpr-question_encoder-single-nq-base",
    "facebook/dpr-question_encoder-multiset-base",
]
TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/dpr-reader-single-nq-base",
    "facebook/dpr-reader-multiset-base",
]







@dataclass
class TFDPRContextEncoderOutput(ModelOutput):
    r

    pooler_output: tf.Tensor = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None


@dataclass
class TFDPRQuestionEncoderOutput(ModelOutput):
    

    pooler_output: tf.Tensor = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None


@dataclass
class TFDPRReaderOutput(ModelOutput):
    

    start_logits: tf.Tensor = None
    end_logits: tf.Tensor = None
    relevance_logits: tf.Tensor = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None


class TFDPREncoderLayer(tf.keras.layers.Layer):
    base_model_prefix = "bert_model"

    def __init__(self, config: DPRConfig, **kwargs):
        super().__init__(**kwargs)

        
        self.bert_model = TFBertMainLayer(config, add_pooling_layer=False, name="bert_model")
        self.config = config

        if self.config.hidden_size <= 0:
            raise ValueError("Encoder hidden_size can't be zero")
        self.projection_dim = config.projection_dim
        if self.projection_dim > 0:
            self.encode_proj = tf.keras.layers.Dense(
                config.projection_dim, kernel_initializer=get_initializer(config.initializer_range), name="encode_proj"
            )

    @unpack_inputs
    def call(
        self,
        input_ids: tf.Tensor = None,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: bool = None,
        output_hidden_states: bool = None,
        return_dict: bool = None,
        training: bool = False,
    ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor, ...]]:
        outputs = self.bert_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        sequence_output = outputs[0]
        pooled_output = sequence_output[:, 0, :]
        if self.projection_dim > 0:
            pooled_output = self.encode_proj(pooled_output)

        if not return_dict:
            return (sequence_output, pooled_output) + outputs[1:]

        return TFBaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    @property
    def embeddings_size(self) -> int:
        if self.projection_dim > 0:
            return self.projection_dim
        return self.bert_model.config.hidden_size


class TFDPRSpanPredictorLayer(tf.keras.layers.Layer):
    base_model_prefix = "encoder"

    def __init__(self, config: DPRConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.encoder = TFDPREncoderLayer(config, name="encoder")

        self.qa_outputs = tf.keras.layers.Dense(
            2, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
        )
        self.qa_classifier = tf.keras.layers.Dense(
            1, kernel_initializer=get_initializer(config.initializer_range), name="qa_classifier"
        )

    @unpack_inputs
    def call(
        self,
        input_ids: tf.Tensor = None,
        attention_mask: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = False,
        training: bool = False,
    ) -> Union[TFDPRReaderOutput, Tuple[tf.Tensor, ...]]:
        
        n_passages, sequence_length = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)[:2]
        
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        sequence_output = outputs[0]

        
        logits = self.qa_outputs(sequence_output)
        start_logits, end_logits = tf.split(logits, 2, axis=-1)
        start_logits = tf.squeeze(start_logits, axis=-1)
        end_logits = tf.squeeze(end_logits, axis=-1)
        relevance_logits = self.qa_classifier(sequence_output[:, 0, :])

        
        start_logits = tf.reshape(start_logits, [n_passages, sequence_length])
        end_logits = tf.reshape(end_logits, [n_passages, sequence_length])
        relevance_logits = tf.reshape(relevance_logits, [n_passages])

        if not return_dict:
            return (start_logits, end_logits, relevance_logits) + outputs[2:]

        return TFDPRReaderOutput(
            start_logits=start_logits,
            end_logits=end_logits,
            relevance_logits=relevance_logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class TFDPRSpanPredictor(TFPreTrainedModel):
    base_model_prefix = "encoder"

    def __init__(self, config: DPRConfig, **kwargs):
        super().__init__(config, **kwargs)
        self.encoder = TFDPRSpanPredictorLayer(config)

    @unpack_inputs
    def call(
        self,
        input_ids: tf.Tensor = None,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = False,
        training: bool = False,
    ) -> Union[TFDPRReaderOutput, Tuple[tf.Tensor, ...]]:
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        return outputs


class TFDPREncoder(TFPreTrainedModel):
    base_model_prefix = "encoder"

    def __init__(self, config: DPRConfig, **kwargs):
        super().__init__(config, **kwargs)

        self.encoder = TFDPREncoderLayer(config)

    @unpack_inputs
    def call(
        self,
        input_ids: tf.Tensor = None,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = False,
        training: bool = False,
    ) -> Union[TFDPRReaderOutput, Tuple[tf.Tensor, ...]]:
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        return outputs







class TFDPRPretrainedContextEncoder(TFPreTrainedModel):
    

    config_class = DPRConfig
    base_model_prefix = "ctx_encoder"


class TFDPRPretrainedQuestionEncoder(TFPreTrainedModel):
    

    config_class = DPRConfig
    base_model_prefix = "question_encoder"


class TFDPRPretrainedReader(TFPreTrainedModel):
    

    config_class = DPRConfig
    base_model_prefix = "reader"







TF_DPR_START_DOCSTRING = r

TF_DPR_ENCODERS_INPUTS_DOCSTRING = r

TF_DPR_READER_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The bare DPRContextEncoder transformer outputting pooler outputs as context representations.",
    TF_DPR_START_DOCSTRING,
)
class TFDPRContextEncoder(TFDPRPretrainedContextEncoder):
    def __init__(self, config: DPRConfig, *args, **kwargs):
        super().__init__(config, *args, **kwargs)
        self.ctx_encoder = TFDPREncoderLayer(config, name="ctx_encoder")

    def get_input_embeddings(self):
        try:
            return self.ctx_encoder.bert_model.get_input_embeddings()
        except AttributeError:
            self.build()
            return self.ctx_encoder.bert_model.get_input_embeddings()

    @unpack_inputs
    @add_start_docstrings_to_model_forward(TF_DPR_ENCODERS_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFDPRContextEncoderOutput, config_class=_CONFIG_FOR_DOC)
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        training: bool = False,
    ) -> TFDPRContextEncoderOutput | Tuple[tf.Tensor, ...]:
        r
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = shape_list(input_ids)
        elif inputs_embeds is not None:
            input_shape = shape_list(inputs_embeds)[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if attention_mask is None:
            attention_mask = (
                tf.ones(input_shape, dtype=tf.dtypes.int32)
                if input_ids is None
                else (input_ids != self.config.pad_token_id)
            )
        if token_type_ids is None:
            token_type_ids = tf.zeros(input_shape, dtype=tf.dtypes.int32)

        outputs = self.ctx_encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        if not return_dict:
            return outputs[1:]

        return TFDPRContextEncoderOutput(
            pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions
        )


@add_start_docstrings(
    "The bare DPRQuestionEncoder transformer outputting pooler outputs as question representations.",
    TF_DPR_START_DOCSTRING,
)
class TFDPRQuestionEncoder(TFDPRPretrainedQuestionEncoder):
    def __init__(self, config: DPRConfig, *args, **kwargs):
        super().__init__(config, *args, **kwargs)
        self.question_encoder = TFDPREncoderLayer(config, name="question_encoder")

    def get_input_embeddings(self):
        try:
            return self.question_encoder.bert_model.get_input_embeddings()
        except AttributeError:
            self.build()
            return self.question_encoder.bert_model.get_input_embeddings()

    @unpack_inputs
    @add_start_docstrings_to_model_forward(TF_DPR_ENCODERS_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFDPRQuestionEncoderOutput, config_class=_CONFIG_FOR_DOC)
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        training: bool = False,
    ) -> TFDPRQuestionEncoderOutput | Tuple[tf.Tensor, ...]:
        r
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = shape_list(input_ids)
        elif inputs_embeds is not None:
            input_shape = shape_list(inputs_embeds)[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if attention_mask is None:
            attention_mask = (
                tf.ones(input_shape, dtype=tf.dtypes.int32)
                if input_ids is None
                else (input_ids != self.config.pad_token_id)
            )
        if token_type_ids is None:
            token_type_ids = tf.zeros(input_shape, dtype=tf.dtypes.int32)

        outputs = self.question_encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        if not return_dict:
            return outputs[1:]
        return TFDPRQuestionEncoderOutput(
            pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions
        )


@add_start_docstrings(
    "The bare DPRReader transformer outputting span predictions.",
    TF_DPR_START_DOCSTRING,
)
class TFDPRReader(TFDPRPretrainedReader):
    def __init__(self, config: DPRConfig, *args, **kwargs):
        super().__init__(config, *args, **kwargs)
        self.span_predictor = TFDPRSpanPredictorLayer(config, name="span_predictor")

    def get_input_embeddings(self):
        try:
            return self.span_predictor.encoder.bert_model.get_input_embeddings()
        except AttributeError:
            self.build()
            return self.span_predictor.encoder.bert_model.get_input_embeddings()

    @unpack_inputs
    @add_start_docstrings_to_model_forward(TF_DPR_READER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFDPRReaderOutput, config_class=_CONFIG_FOR_DOC)
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        training: bool = False,
    ) -> TFDPRReaderOutput | Tuple[tf.Tensor, ...]:
        r
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = shape_list(input_ids)
        elif inputs_embeds is not None:
            input_shape = shape_list(inputs_embeds)[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if attention_mask is None:
            attention_mask = tf.ones(input_shape, dtype=tf.dtypes.int32)

        return self.span_predictor(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )