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