Upload model

bert_classification_model.py

@@ -4,8 +4,11 @@ from typing import Tuple
 from transformers import PreTrainedModel, BertModel
 from torch import nn
 
+from bert_classification_config import BertClassificationConfig
+
 
 class BertClassificationModel(PreTrainedModel):
+    config_class = BertClassificationConfig
 
     def __init__(self, config, num_main_segment=None, num_sub_segment=None):
         super(BertClassificationModel, self).__init__(config=config)

config.json

@@ -14,7 +14,7 @@
   "intermediate_size": 3072,
   "layer_norm_eps": 1e-12,
   "max_position_embeddings": 512,
-  "model_type": "bert",
+  "model_type": "bert-classification",
   "num_attention_heads": 12,
   "num_hidden_layers": 12,
   "num_main_segment": 0,

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:18b607e5e17d75afbb4b65e8c03fda9e28d87a1c9259a4c9cc5f291696317003
+oid sha256:fc4263b34d694e1985ba5f30ff508f1d307614e09914270b24136661769254c9
 size 671857125