| """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: |
| |
| 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) |
| |
| 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") |
| |
| 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)))) |
| |
| 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()) |
|
|
