Upload BERTClassifier

config.json

@@ -4,6 +4,9 @@
     "BERTClassifier"
   ],
   "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoModel": "model.BERTClassifier"
+  },
   "classifier_dropout": null,
   "gradient_checkpointing": false,
   "hidden_act": "gelu",

model.py

@@ -0,0 +1,28 @@
+from torch import nn
+from transformers import BertModel, PreTrainedModel
+
+class BERTClassifier(PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = BertModel(config)
+        self.dropout = nn.Dropout(0.1)
+        self.fc = nn.Linear(self.bert.config.hidden_size, 16)
+
+    def forward(self, input_ids, attention_mask):
+            outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+            pooled_output = outputs.pooler_output
+            x = self.dropout(pooled_output)
+            logits = self.fc(x)
+            return logits
+    
+    def print_test(self, x):
+        return "lmao"
+
+
+if __name__ == "__main__":
+
+    from transformers import BertConfig, BertModel, BertForMaskedLM, AutoConfig
+
+    # Initializing a BERT google-bert/bert-base-uncased style configuration
+    config = AutoConfig.from_pretrained('google-bert/bert-base-uncased', trust_remote_code=True)
+    model = BERTClassifier(config)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f135923bede3e2e2642a25d8dd463d90909ecdd748dd82fff62adbfe2a844bd0
+oid sha256:a7bd7ed63579a02de09de21fe7d95b4df37dcb129ecef34c95653af627cc5745
 size 438001696