Upload PopBERT

config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "PopBERT"
   ],
+  "auto_map": {
+    "AutoConfig": "module.PopBERTConfig",
+    "AutoModel": "module.PopBERT"
+  },
   "model_type": "popbert",
   "num_classes": 4,
   "torch_dtype": "float32",

module.py

@@ -0,0 +1,52 @@
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss
+from transformers import AutoModelForSequenceClassification
+from transformers import PretrainedConfig
+from transformers import PreTrainedModel
+
+
+class PopBERTConfig(PretrainedConfig):
+    model_type = "popbert"
+
+    def __init__(self, num_classes: int = 4, **kwargs):
+        super().__init__(**kwargs)
+        self.num_classes = num_classes
+
+
+class PopBERT(PreTrainedModel):
+    config_class = PopBERTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.sigmoid = nn.Sigmoid()
+        self.bert = AutoModelForSequenceClassification.from_pretrained(
+            "deepset/gbert-large",
+            num_labels=config.num_classes,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        head_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+    ):
+        pred = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+        )
+
+        loss = None
+        if labels is not None:
+            loss_fn = BCEWithLogitsLoss()
+            loss = loss_fn(pred.logits, labels.float())
+
+        return loss, self.sigmoid(pred.logits)

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:17b602ee2fa8553d861f68a2420e5c136b6dca34b5f77a5939417ab075feb06e
+oid sha256:6483ab1821245614459ff691f00ad8dd291e37e39c6952c3e67b5427f2c6ea22
 size 1343094965