bert_base_pretrain

→BERTForPreTraining

bert_base_pretrain(pretrained: bool | str = False, weights: BERTBaseWeights | None = None, overrides: object = {})

source edit

Construct BERT-Base with both pre-training heads attached.

The masked-LM head and the next-sentence head at once, which is the configuration the paper trains: the two losses are summed, and the forward returns each separately so a caller can weight or inspect them. This is the factory for continued pre-training on a domain corpus; bert_base_mlm is the one for mask-filling alone.

Model Size

Parameters

pretrainedbool or str= False

Encoder-weight selector. False → fully random init; True → loads the bert_base encoder (BERTBaseWeights.DEFAULT) into the .bert trunk. Both heads are randomly initialised.

weights(BERTBaseWeights, optional, keyword - only)= None

Explicit encoder-weights enum member; takes precedence over pretrained.

**overridesobject= {}

Optional BERTConfig field overrides.

Returns

BERTForPreTraining

BERT-Base wrapped with the masked-LM and next-sentence heads.

Notes

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

The output carries prediction_logits and seq_relationship_logits rather than a single logits — two heads produce two predictions, and mlm_loss / nsp_loss are reported beside the summed loss so the balance between them is visible.

Examples

>>> import lucid
>>> from lucid.models.text.bert import bert_base_pretrain
>>> model = bert_base_pretrain(vocab_size=99, hidden_size=32,
...                            num_hidden_layers=2,
...                            num_attention_heads=2,
...                            intermediate_size=64).eval()
>>> out = model(lucid.tensor([[1, 2, 3, 4]]))
>>> out.prediction_logits.shape, out.seq_relationship_logits.shape
((1, 4, 99), (1, 2))

Used by 1

lucid.models.text.bert