add base model option

model.py

@@ -4,10 +4,14 @@ from transformers.modeling_outputs import TokenClassifierOutput
 
 
 class BertClassifier(nn.Module):
-    def __init__(self, num_labels=2, dropout=0.1):
+    def __init__(self, num_labels=2, dropout=0.1,bert_model=None):
         super().__init__()
-        config = BertConfig(vocab_size=34688, max_position_embeddings=512)
-        self.bert = BertModel(config=config)
+        if bert_model:
+            self.bert = BertModel.from_pretrained(bert_model)
+        else:
+            config = BertConfig(vocab_size=34688, max_position_embeddings=512)
+            self.bert = BertModel(config=config)
+            
         self.num_labels = num_labels
         self.classifier = nn.Sequential(
             nn.Linear(self.bert.config.hidden_size, self.bert.config.hidden_size),