class

BERTForQuestionAnswering

extendsSequenceClassificationModel

BERTForQuestionAnswering(config: BERTConfig)

source edit

BERT with a 2-way span head for extractive question answering.

Wraps the bidirectional encoder with a single linear of output width 2, producing start- and end-position logits over each token in the input. This is the SQuAD v1.1 / v2.0 fine-tuning recipe of Devlin et al., 2018 §4.2 — given a (question, context) pair concatenated with [SEP], the model predicts the answer span inside the context.

Parameters

configBERTConfig

BERT hyperparameters. The QA head is always 2-way; num_labels is ignored here.

Attributes

bertBERTModel

Underlying bidirectional encoder trunk.

qa_outputsnn.Linear

Final linear of shape (hidden_size, 2) mapping each token's hidden state to (start_logit, end_logit).

Notes

Reference: Devlin, Chang, Lee, and Toutanova, "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding", NAACL 2019 (arXiv:1810.04805).

When both start_positions and end_positions are provided, the loss is the symmetric average of two cross-entropies:

\mathcal{L} = \tfrac{1}{2}\!\left( \mathrm{CE}(z^{\mathrm{start}}, y^{\mathrm{start}}) + \mathrm{CE}(z^{\mathrm{end}}, y^{\mathrm{end}}) \right).

The forward returns a QuestionAnsweringOutput carrying start_logits and end_logits, each of shape (B, T).

Examples

>>> import lucid
>>> from lucid.models.text.bert import BERTConfig, BERTForQuestionAnswering
>>> cfg = BERTConfig(num_hidden_layers=2, hidden_size=128,
...                  num_attention_heads=2, intermediate_size=512)
>>> model = BERTForQuestionAnswering(cfg).eval()
>>> input_ids = lucid.tensor([[101, 2040, 2003, 102, 1045, 2572, 102]])
>>> out = model(input_ids)
>>> out.start_logits.shape, out.end_logits.shape
((1, 7), (1, 7))

Used by 2

Constructors

dunder

init

→None

__init__(config: BERTConfig)

source edit

Instance methods

forward

→QuestionAnsweringOutput

forward(input_ids: Tensor, attention_mask: Tensor | None = None, token_type_ids: Tensor | None = None, start_positions: Tensor | None = None, end_positions: Tensor | None = None)

source edit

Sum the word, position and token-type embeddings.

Exactly one of input_ids and inputs_embeds is required. inputs_embeds skips the word-embedding lookup, which is what a caller doing soft prompts or adapter tuning needs; position_ids overrides the default 0..T-1, which packed or shifted sequences need.

>>> import lucid >>> from lucid.models.text.bert import BERTConfig, BERTForQuestionAnswering >>> cfg = BERTConfig(num_hidden_layers=2, hidden_size=128, ... num_attention_heads=2, intermediate_size=512) >>> model = BERTForQuestionAnswering(cfg).eval() >>> input_ids = lucid.tensor([[101, 2040, 2003, 102, 1045, 2572, 102]]) >>> out = model(input_ids) >>> out.start_logits.shape, out.end_logits.shape ((1, 7), (1, 7))