Update geneformer/mtl/model.py

change in line 36:
old:
self.bert = BertModel(self.config)
new:
self.bert = BertModel.from_pretrained(pretrained_path)

Reason:
BertModel() is initializing random weights instead of using weights from pretrained model. This results in training from scratch instead of finetuning the model.
BertModel.from_pretrained() initializes the pretrained weights.
Tested with training with MTLClassifier both ways and also via:
pretrained_path = "/Geneformer/Geneformer-V2-104M"
ref = BertModel.from_pretrained(pretrained_path)
mtl = GeneformerMultiTask(pretrained_path, num_labels_list=[2])

print("mtl layer0 q mean:", mtl.bert.encoder.layer[0].attention.self.query.weight.abs().mean().item())
print("ref layer0 q mean:", ref.encoder.layer[0].attention.self.query.weight.abs().mean().item())

geneformer/mtl/model.py

@@ -33,7 +33,7 @@ class GeneformerMultiTask(nn.Module):
     ):
         super(GeneformerMultiTask, self).__init__()
         self.config = BertConfig.from_pretrained(pretrained_path)
-        self.bert = BertModel(self.config)
+        self.bert = BertModel.from_pretrained(pretrained_path)
         self.num_labels_list = num_labels_list
         self.use_task_weights = use_task_weights
         self.dropout = nn.Dropout(dropout_rate)