class

TransformerForTokenClassification

extendsPretrainedModelMaskedLMMixin

TransformerForTokenClassification(config: TransformerConfig)

source edit

Encoder-only per-token classifier — NER / POS / chunking heads.

Wraps TransformerModel and runs only the encoder half during forward; every source-side token's hidden state is dropout-regularised and projected through a linear of output width config.num_labels. Uses MaskedLMMixin for the standard masked (B, T, V) → loss reduction shared with BERT and RoFormer.

Parameters

configTransformerConfig

Hyperparameters. config.num_labels sets the per-position output dimension; config.classifier_dropout (falling back to hidden_dropout) sets the dropout applied before the linear.

Attributes

transformerTransformerModel

Underlying encoder-decoder trunk; only encode is invoked here.

dropoutnn.Dropout

Dropout layer applied to the full sequence hidden states.

classifiernn.Linear

Final linear of shape (d_model, num_labels) mapping each token's hidden state to per-class logits.

Notes

Reference: Vaswani et al., "Attention Is All You Need", NeurIPS, 2017 (arXiv:1706.03762).

Loss (when labels is provided) is the masked cross-entropy over positions whose label is not -100:

\mathcal{L} = -\frac{1}{|V|} \sum_{(b, t) \in V} \log p_\theta(y_{b, t} \mid x_b).

Examples

>>> import lucid
>>> from lucid.models.text.transformer import (
...     TransformerConfig, TransformerForTokenClassification,
... )
>>> cfg = TransformerConfig(num_labels=9, num_hidden_layers=2,
...                         hidden_size=64, num_attention_heads=2,
...                         intermediate_size=128, vocab_size=1000)
>>> model = TransformerForTokenClassification(cfg).eval()
>>> input_ids = lucid.tensor([[1, 234, 567, 2]])
>>> out = model(input_ids)
>>> out.logits.shape   # (B=1, T=4, num_labels=9)
(1, 4, 9)

Used by 2

Constructors

dunder

init

→None

__init__(config: TransformerConfig)

source edit

Instance methods

forward

→MaskedLMOutput

forward(input_ids: Tensor, attention_mask: Tensor | None = None, labels: Tensor | None = None)

source edit

Full encode + decode pass.

Parameters

input_ids

(B, S) source token ids.

decoder_input_ids

(B, T) target token ids (right-shifted).

attention_mask

Optional (B, S) HF-style padding mask.

decoder_attention_mask

Optional (B, T) HF-style padding mask.

Returns

MaskedLMOutput

class:Seq2SeqLMOutput carrying decoder hidden states as

>>> import lucid >>> from lucid.models.text.transformer import ( ... TransformerConfig, TransformerForTokenClassification, ... ) >>> cfg = TransformerConfig(num_labels=9, num_hidden_layers=2, ... hidden_size=64, num_attention_heads=2, ... intermediate_size=128, vocab_size=1000) >>> model = TransformerForTokenClassification(cfg).eval() >>> input_ids = lucid.tensor([[1, 234, 567, 2]]) >>> out = model(input_ids) >>> out.logits.shape # (B=1, T=4, num_labels=9) (1, 4, 9)