Upload ProkBertForMaskedLM

config.json

@@ -5,9 +5,9 @@
   ],
   "attention_probs_dropout_prob": 0.1,
   "auto_map": {
-    "AutoConfig": "neuralbioinfo/prokbert-mini-long--models.ProkBertConfig",
+    "AutoConfig": "models.ProkBertConfig",
     "AutoModel": "neuralbioinfo/prokbert-mini-long--models.ProkBertModel",
-    "AutoModelForMaskedLM": "neuralbioinfo/prokbert-mini-long--models.ProkBertForMaskedLM",
+    "AutoModelForMaskedLM": "models.ProkBertForMaskedLM",
     "AutoModelForSequenceClassification": "neuralbioinfo/prokbert-mini-long--models.ProkBertForSequenceClassification"
   },
   "classification_dropout_rate": 0.1,

models.py

@@ -9,7 +9,7 @@ import torch.nn.functional as F
 from transformers import MegatronBertConfig, MegatronBertModel, MegatronBertForMaskedLM, MegatronBertPreTrainedModel, PreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput
 from transformers.utils.hub import cached_file
-
+from prokbert.training_utils import compute_metrics_eval_prediction
 
 class BertForBinaryClassificationWithPooling(nn.Module):
     """
@@ -130,18 +130,6 @@ class BertForBinaryClassificationWithPooling(nn.Module):
 
 
 
-class OldProkBertConfig(MegatronBertConfig):
-
-    model_type = "prokbert"
-    def __init__(
-        self,
-        kmer: int = 6, 
-        shift: int = 1,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.kmer=kmer
-        self.shift=shift
 
 class ProkBertConfig(MegatronBertConfig):
     model_type = "prokbert"
@@ -283,8 +271,10 @@ class ProkBertForSequenceClassification(ProkBertPreTrainedModel):
             # Classification head
             pooled_output = self.dropout(pooled_output)
             logits = self.classifier(pooled_output)
-            loss = self.loss_fct(logits.view(-1, 2), labels.view(-1))
-
+            loss = None
+            if labels is not None:
+                loss = self.loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
+                
             classification_output = SequenceClassifierOutput(
                 loss=loss,
                 logits=logits,