class

RoFormerForQuestionAnswering

extendsSequenceClassificationModel

RoFormerForQuestionAnswering(config: RoFormerConfig)

source edit

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

Wraps RoFormerModel with a single linear of output width 2, producing start- and end-position logits over each token in the input. Identical contract to BERTForQuestionAnswering; the SQuAD v1.1 / v2.0 fine-tuning recipe applied to a RoPE-trained encoder.

Parameters

configRoFormerConfig

Hyperparameters. The QA head is always 2-way; num_labels is ignored.

Attributes

roformerRoFormerModel

Underlying RoPE 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: Su, Lu, Pan, Murtadha, Wen, and Liu, "RoFormer: Enhanced Transformer with Rotary Position Embedding", Neurocomputing, vol. 568, 2024 (arXiv:2104.09864).

When both start_positions and end_positions are supplied 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 two score vectors are returned separately as start_logits and end_logits, each (B, T), in a lucid.models.QuestionAnsweringOutput — the same shape lucid.models.BERTForQuestionAnswering returns. It used to hand back one (B, T, 2) tensor and leave the caller to remember which index was which, which made the "identical contract" claimed above false for the one thing a caller actually touches.

Examples

>>> import lucid
>>> from lucid.models.text.roformer import (
...     RoFormerConfig, RoFormerForQuestionAnswering,
... )
>>> cfg = RoFormerConfig(num_hidden_layers=2, hidden_size=128,
...                      num_attention_heads=2, intermediate_size=512)
>>> model = RoFormerForQuestionAnswering(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: RoFormerConfig)

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

>>> import lucid >>> from lucid.models.text.roformer import ( ... RoFormerConfig, RoFormerForQuestionAnswering, ... ) >>> cfg = RoFormerConfig(num_hidden_layers=2, hidden_size=128, ... num_attention_heads=2, intermediate_size=512) >>> model = RoFormerForQuestionAnswering(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))