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 pytorch dpr model for open domain question answering from dataclasses import dataclass from typing import optional tuple union import torch from torch import tensor nn from modelingoutputs import basemodeloutputwithpooling from modelingutils import pretrainedmodel from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from bert modelingbert import bertmodel from configurationdpr import dprconfig logger logging getloggername configfordoc dprconfig checkpointfordoc facebookdprctxencodersinglenqbase dprcontextencoderpretrainedmodelarchivelist facebookdprctxencodersinglenqbase facebookdprctxencodermultisetbase dprquestionencoderpretrainedmodelarchivelist facebookdprquestionencodersinglenqbase facebookdprquestionencodermultisetbase dprreaderpretrainedmodelarchivelist facebookdprreadersinglenqbase facebookdprreadermultisetbase outputs dataclass class dprcontextencoderoutputmodeloutput pooleroutput torch floattensor hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none dataclass class dprquestionencoderoutputmodeloutput pooleroutput torch floattensor hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none dataclass class dprreaderoutputmodeloutput startlogits torch floattensor endlogits torch floattensor none relevancelogits torch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none class dprpretrainedmodelpretrainedmodel def initweightsself module slightly different from the tf version which uses truncatednormal for initialization cf https github compytorchpytorchpull5617 initialize weights and apply final processing initialize weights and apply final processing 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 dprencodersinputsdocstring r args inputids torch longtensor 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 torch floattensor 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 torch longtensor 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 torch floattensor 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 outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple addstartdocstrings the bare dprcontextencoder transformer outputting pooler outputs as context representations dprstartdocstring class dprcontextencoderdprpretrainedcontextencoder def initself config dprconfig super initconfig self config config self ctxencoder dprencoderconfig initialize weights and apply final processing self postinit addstartdocstringstomodelforwarddprencodersinputsdocstring replacereturndocstringsoutputtypedprcontextencoderoutput configclassconfigfordoc def forward self inputids optionaltensor none attentionmask optionaltensor none tokentypeids optionaltensor none inputsembeds optionaltensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniondprcontextencoderoutput tupletensor r return examples python from transformers import dprcontextencoder dprcontextencodertokenizer tokenizer dprcontextencodertokenizer frompretrainedfacebookdprctxencodersinglenqbase model dprcontextencoder frompretrainedfacebookdprctxencodersinglenqbase inputids tokenizerhello is my dog cute returntensorsptinputids embeddings modelinputids pooleroutput outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict 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 inputids size elif inputsembeds is not none inputshape inputsembeds size 1 else raise valueerroryou have to specify either inputids or inputsembeds device inputids device if inputids is not none else inputsembeds device if attentionmask is none attentionmask torch onesinputshape devicedevice if inputids is none else inputids self config padtokenid if tokentypeids is none tokentypeids torch zerosinputshape dtypetorch long devicedevice outputs self ctxencoder inputidsinputids attentionmaskattentionmask tokentypeidstokentypeids inputsembedsinputsembeds outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict if not returndict return outputs1 return dprcontextencoderoutput pooleroutputoutputs pooleroutput hiddenstatesoutputs hiddenstates attentionsoutputs attentions addstartdocstrings the bare dprquestionencoder transformer outputting pooler outputs as question representations dprstartdocstring class dprquestionencoderdprpretrainedquestionencoder def initself config dprconfig super initconfig self config config self questionencoder dprencoderconfig initialize weights and apply final processing self postinit addstartdocstringstomodelforwarddprencodersinputsdocstring replacereturndocstringsoutputtypedprquestionencoderoutput configclassconfigfordoc def forward self inputids optionaltensor none attentionmask optionaltensor none tokentypeids optionaltensor none inputsembeds optionaltensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniondprquestionencoderoutput tupletensor r return examples python from transformers import dprquestionencoder dprquestionencodertokenizer tokenizer dprquestionencodertokenizer frompretrainedfacebookdprquestionencodersinglenqbase model dprquestionencoder frompretrainedfacebookdprquestionencodersinglenqbase inputids tokenizerhello is my dog cute returntensorsptinputids embeddings modelinputids pooleroutput initialize weights and apply final processing outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict 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 self warnifpaddingandnoattentionmaskinputids attentionmask inputshape inputids size elif inputsembeds is not none inputshape inputsembeds size 1 else raise valueerroryou have to specify either inputids or inputsembeds device inputids device if inputids is not none else inputsembeds device if attentionmask is none attentionmask torch onesinputshape devicedevice return self spanpredictor inputids attentionmask inputsembedsinputsembeds outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict 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 pytorch dpr model for open domain question answering outputs class for outputs of dprquestionencoder args pooler_output torch floattensor 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 torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor 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 class for outputs of dprquestionencoder args pooler_output torch floattensor 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 torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor 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 class for outputs of dprquestionencoder args start_logits torch floattensor of shape n_passages sequence_length logits of the start index of the span for each passage end_logits torch floattensor of shape n_passages sequence_length logits of the end index of the span for each passage relevance_logits torch floattensor 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 torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor 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 initialize the weights slightly different from the tf version which uses truncated_normal for initialization cf https github com pytorch pytorch pull 5617 initialize weights and apply final processing initialize weights and apply final processing 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 pretrainedmodel 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 pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior 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 pretrainedmodel from_pretrained method to load the model weights args input_ids torch longtensor 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 torch floattensor 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 torch longtensor 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 torch floattensor 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 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 args input_ids tuple torch longtensor 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 what are input ids glossary input ids attention_mask torch floattensor 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 torch floattensor 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_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail 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 initialize weights and apply final processing return examples python from transformers import dprcontextencoder dprcontextencodertokenizer tokenizer dprcontextencodertokenizer from_pretrained facebook dpr ctx_encoder single nq base model dprcontextencoder from_pretrained facebook dpr ctx_encoder single nq base input_ids tokenizer hello is my dog cute return_tensors pt input_ids embeddings model input_ids pooler_output initialize weights and apply final processing return examples python from transformers import dprquestionencoder dprquestionencodertokenizer tokenizer dprquestionencodertokenizer from_pretrained facebook dpr question_encoder single nq base model dprquestionencoder from_pretrained facebook dpr question_encoder single nq base input_ids tokenizer hello is my dog cute return_tensors pt input_ids embeddings model input_ids pooler_output initialize weights and apply final processing return examples python from transformers import dprreader dprreadertokenizer tokenizer dprreadertokenizer from_pretrained facebook dpr reader single nq base model dprreader from_pretrained facebook dpr reader single nq base encoded_inputs tokenizer questions what is love titles haddaway texts what is love is a song recorded by the artist haddaway return_tensors pt outputs model encoded_inputs start_logits outputs start_logits end_logits outputs end_logits relevance_logits outputs relevance_logits
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
from torch import Tensor, nn

from ...modeling_outputs import BaseModelOutputWithPooling
from ...modeling_utils import PreTrainedModel
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ..bert.modeling_bert import BertModel
from .configuration_dpr import DPRConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "DPRConfig"
_CHECKPOINT_FOR_DOC = "facebook/dpr-ctx_encoder-single-nq-base"

DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/dpr-ctx_encoder-single-nq-base",
    "facebook/dpr-ctx_encoder-multiset-base",
]
DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/dpr-question_encoder-single-nq-base",
    "facebook/dpr-question_encoder-multiset-base",
]
DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/dpr-reader-single-nq-base",
    "facebook/dpr-reader-multiset-base",
]







@dataclass
class DPRContextEncoderOutput(ModelOutput):
    

    pooler_output: torch.FloatTensor
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class DPRQuestionEncoderOutput(ModelOutput):
    

    pooler_output: torch.FloatTensor
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class DPRReaderOutput(ModelOutput):
    

    start_logits: torch.FloatTensor
    end_logits: torch.FloatTensor = None
    relevance_logits: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


class DPRPreTrainedModel(PreTrainedModel):
    def _init_weights(self, module):
        
        if isinstance(module, nn.Linear):
            
            
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


class DPREncoder(DPRPreTrainedModel):
    base_model_prefix = "bert_model"

    def __init__(self, config: DPRConfig):
        super().__init__(config)
        self.bert_model = BertModel(config, add_pooling_layer=False)
        if self.bert_model.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 = nn.Linear(self.bert_model.config.hidden_size, config.projection_dim)
        
        self.post_init()

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Optional[Tensor] = None,
        token_type_ids: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = False,
    ) -> Union[BaseModelOutputWithPooling, Tuple[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,
        )
        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[2:]

        return BaseModelOutputWithPooling(
            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.encode_proj.out_features
        return self.bert_model.config.hidden_size


class DPRSpanPredictor(DPRPreTrainedModel):
    base_model_prefix = "encoder"

    def __init__(self, config: DPRConfig):
        super().__init__(config)
        self.encoder = DPREncoder(config)
        self.qa_outputs = nn.Linear(self.encoder.embeddings_size, 2)
        self.qa_classifier = nn.Linear(self.encoder.embeddings_size, 1)
        
        self.post_init()

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor,
        inputs_embeds: Optional[Tensor] = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = False,
    ) -> Union[DPRReaderOutput, Tuple[Tensor, ...]]:
        
        n_passages, sequence_length = input_ids.size() if input_ids is not None else inputs_embeds.size()[:2]
        
        outputs = self.encoder(
            input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = outputs[0]

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

        
        start_logits = start_logits.view(n_passages, sequence_length)
        end_logits = end_logits.view(n_passages, sequence_length)
        relevance_logits = relevance_logits.view(n_passages)

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

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







class DPRPretrainedContextEncoder(DPRPreTrainedModel):
    

    config_class = DPRConfig
    load_tf_weights = None
    base_model_prefix = "ctx_encoder"


class DPRPretrainedQuestionEncoder(DPRPreTrainedModel):
    

    config_class = DPRConfig
    load_tf_weights = None
    base_model_prefix = "question_encoder"


class DPRPretrainedReader(DPRPreTrainedModel):
    

    config_class = DPRConfig
    load_tf_weights = None
    base_model_prefix = "span_predictor"







DPR_START_DOCSTRING = r

DPR_ENCODERS_INPUTS_DOCSTRING = r

DPR_READER_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The bare DPRContextEncoder transformer outputting pooler outputs as context representations.",
    DPR_START_DOCSTRING,
)
class DPRContextEncoder(DPRPretrainedContextEncoder):
    def __init__(self, config: DPRConfig):
        super().__init__(config)
        self.config = config
        self.ctx_encoder = DPREncoder(config)
        
        self.post_init()

    @add_start_docstrings_to_model_forward(DPR_ENCODERS_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=DPRContextEncoderOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        token_type_ids: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[DPRContextEncoderOutput, Tuple[Tensor, ...]]:
        r

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

        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 = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
            attention_mask = (
                torch.ones(input_shape, device=device)
                if input_ids is None
                else (input_ids != self.config.pad_token_id)
            )
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        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,
        )

        if not return_dict:
            return outputs[1:]
        return DPRContextEncoderOutput(
            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.",
    DPR_START_DOCSTRING,
)
class DPRQuestionEncoder(DPRPretrainedQuestionEncoder):
    def __init__(self, config: DPRConfig):
        super().__init__(config)
        self.config = config
        self.question_encoder = DPREncoder(config)
        
        self.post_init()

    @add_start_docstrings_to_model_forward(DPR_ENCODERS_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=DPRQuestionEncoderOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        token_type_ids: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[DPRQuestionEncoderOutput, Tuple[Tensor, ...]]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        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:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
            attention_mask = (
                torch.ones(input_shape, device=device)
                if input_ids is None
                else (input_ids != self.config.pad_token_id)
            )
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        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,
        )

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


@add_start_docstrings(
    "The bare DPRReader transformer outputting span predictions.",
    DPR_START_DOCSTRING,
)
class DPRReader(DPRPretrainedReader):
    def __init__(self, config: DPRConfig):
        super().__init__(config)
        self.config = config
        self.span_predictor = DPRSpanPredictor(config)
        
        self.post_init()

    @add_start_docstrings_to_model_forward(DPR_READER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=DPRReaderOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[DPRReaderOutput, Tuple[Tensor, ...]]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        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:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
            attention_mask = torch.ones(input_shape, device=device)

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