File size: 6,966 Bytes
90316bb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """Original BERT heads/masking and explicit ModernBERT comparison variants."""
import torch
from torch import nn
from transformers import BertConfig, BertModel, ModernBertConfig, ModernBertModel, ModernBertForSequenceClassification
def backbone_config_from_dict(values):
kind = values.get("model_type", "bert")
classes = {"bert": BertConfig, "modernbert": ModernBertConfig}
if kind not in classes:
raise ValueError(f"Unsupported checkpoint backbone: {kind}")
return classes[kind].from_dict(values)
class VisualRelatednessModel(nn.Module):
def __init__(self, bert_config, architecture="bert-cnn", seq_len=64,
cnn_filters=32, cnn_kernel_size=3, num_bert_layers=None,
tune_embeddings=None, backbone=None):
super().__init__()
if architecture not in {"bert", "bert-cnn", "modernbert", "modernbert-cnn"}:
raise ValueError("Unsupported model architecture")
modern = architecture.startswith("modernbert")
if modern != (bert_config.model_type == "modernbert"):
raise ValueError("Architecture and pretrained backbone type do not match")
if not 3 <= seq_len <= bert_config.max_position_embeddings:
raise ValueError("seq_len must be between 3 and BERT's positional embedding limit")
layers = bert_config.num_hidden_layers if num_bert_layers is None else num_bert_layers
if not 0 <= layers <= bert_config.num_hidden_layers:
raise ValueError("num_bert_layers is outside this BERT model's layer range")
if cnn_filters < 1 or cnn_kernel_size < 1:
raise ValueError("CNN filter count and kernel size must be positive")
self.options = dict(architecture=architecture, seq_len=seq_len,
cnn_filters=cnn_filters, cnn_kernel_size=cnn_kernel_size,
num_bert_layers=layers,
tune_embeddings=architecture in {"bert", "modernbert"} if tune_embeddings is None else tune_embeddings)
if modern:
# Let the outer training loop control compilation explicitly.
bert_config.reference_compile = False
bert_config._attn_implementation = "sdpa"
model_class = ModernBertForSequenceClassification if architecture == "modernbert" else ModernBertModel
self.bert = backbone if backbone is not None else model_class(bert_config)
else:
self.bert = backbone if backbone is not None else BertModel(
bert_config, add_pooling_layer=architecture == "bert")
hidden = bert_config.hidden_size
if architecture == "bert":
self.dropout = nn.Dropout(.1)
self.classifier = nn.Linear(hidden, 2)
nn.init.normal_(self.classifier.weight, std=.02)
nn.init.zeros_(self.classifier.bias)
elif architecture != "modernbert":
pooled_length = (seq_len - cnn_kernel_size + 1) // 2
if pooled_length < 1:
raise ValueError("seq_len must be at least cnn_kernel_size + 1")
self.conv = nn.Conv1d(hidden, cnn_filters, cnn_kernel_size)
self.pool = nn.MaxPool1d(2, stride=2)
self.classifier = nn.Linear(pooled_length * cnn_filters, 1)
# Keras Glorot initialization for the original CNN and dense head.
nn.init.xavier_uniform_(self.conv.weight)
nn.init.zeros_(self.conv.bias)
nn.init.xavier_uniform_(self.classifier.weight)
nn.init.zeros_(self.classifier.bias)
self.configure_finetuning()
@classmethod
def from_pretrained(cls, source, revision=None, **options):
architecture = options.get("architecture", "bert-cnn")
if architecture.startswith("modernbert"):
model_class = ModernBertForSequenceClassification if architecture == "modernbert" else ModernBertModel
extra = {"num_labels": 2} if architecture == "modernbert" else {}
bert = model_class.from_pretrained(source, revision=revision, reference_compile=False,
attn_implementation="sdpa", **extra)
else:
bert = BertModel.from_pretrained(source, revision=revision, add_pooling_layer=architecture == "bert")
return cls(bert.config, backbone=bert, **options)
def configure_finetuning(self):
self.bert.requires_grad_(False)
modern = self.options["architecture"].startswith("modernbert")
encoder = self.bert.model if self.options["architecture"] == "modernbert" else self.bert
if self.options["tune_embeddings"]:
encoder.embeddings.requires_grad_(True)
count = self.options["num_bert_layers"]
if count:
layers = encoder.layers if modern else encoder.encoder.layer
for layer in layers[-count:]:
layer.requires_grad_(True)
if modern:
encoder.final_norm.requires_grad_(True)
if self.options["architecture"] == "modernbert":
self.bert.head.requires_grad_(True)
self.bert.classifier.requires_grad_(True)
if self.options["architecture"] == "bert" and self.bert.pooler is not None:
self.bert.pooler.requires_grad_(True)
def forward(self, input_ids, attention_mask, token_type_ids=None):
if self.options["architecture"].startswith("modernbert"):
output = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
if self.options["architecture"] == "modernbert":
return output.logits
else:
output = self.bert(input_ids=input_ids, attention_mask=attention_mask,
token_type_ids=token_type_ids, return_dict=True)
if self.options["architecture"] == "bert":
return self.classifier(self.dropout(output.pooler_output))
if input_ids.shape[1] != self.options["seq_len"]:
raise ValueError("BERT-CNN requires fixed padding to the checkpoint's seq_len")
# The old BertLayer zeroed padded token vectors BEFORE the convolution.
sequence = output.last_hidden_state * attention_mask.unsqueeze(-1).to(output.last_hidden_state.dtype)
pooled = self.pool(torch.relu(self.conv(sequence.transpose(1, 2))))
# Keras Flatten traverses time then filters, rather than filters then time.
flat = pooled.transpose(1, 2).contiguous().flatten(1)
return self.classifier(flat).squeeze(-1)
def probabilities(self, logits):
return logits.softmax(-1)[:, 1] if self.options["architecture"] in {"bert", "modernbert"} else logits.sigmoid()
def loss(self, logits, labels):
if self.options["architecture"] in {"bert", "modernbert"}:
return nn.functional.cross_entropy(logits, labels.long())
return nn.functional.binary_cross_entropy_with_logits(logits, labels.float())
|