Publish ModernBERT-CNN with concise usage, training details and own evaluation metrics
90316bb verified | """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() | |
| 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()) | |