class

BERTForTokenClassification

extendsPretrainedModelMaskedLMMixin

BERTForTokenClassification(config: BERTConfig)

source edit

BERT with a per-token linear classifier for tagging tasks.

Wraps the bidirectional encoder with a dropout-regularised linear head applied independently at every sequence position. Used for token-level fine-tunes such as named-entity recognition (CoNLL-2003), part-of-speech tagging, and chunking — see Devlin et al., 2018 §4.3.

Parameters

configBERTConfig

BERT 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

bertBERTModel

Underlying bidirectional encoder trunk.

dropoutnn.Dropout

Dropout applied to the full sequence hidden states.

classifiernn.Linear

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

Notes

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

The loss (when labels is provided) is the masked cross-entropy inherited from MaskedLMMixin:

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

where $V$ is the set of positions with label != -100.

Examples

>>> import lucid
>>> from lucid.models.text.bert import BERTConfig, BERTForTokenClassification
>>> cfg = BERTConfig(num_labels=9, num_hidden_layers=2, hidden_size=128,
...                  num_attention_heads=2, intermediate_size=512)
>>> model = BERTForTokenClassification(cfg).eval()
>>> input_ids = lucid.tensor([[101, 7592, 2088, 102]])
>>> 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: BERTConfig)

source edit

Instance methods

forward

→MaskedLMOutput

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

source edit

>>> import lucid >>> from lucid.models.text.bert import BERTConfig, BERTForTokenClassification >>> cfg = BERTConfig(num_labels=9, num_hidden_layers=2, hidden_size=128, ... num_attention_heads=2, intermediate_size=512) >>> model = BERTForTokenClassification(cfg).eval() >>> input_ids = lucid.tensor([[101, 7592, 2088, 102]]) >>> out = model(input_ids) >>> out.logits.shape # (B=1, T=4, num_labels=9) (1, 4, 9)