Update modeling_re_bert.py

modeling_re_bert.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+from typing import List, Dict
+
 import json
 import torch
 import torch.nn as nn
@@ -27,14 +29,13 @@ FEW_REL_DEFINITIONS = [
 
 class BertForRelationClassification(BertPreTrainedModel):
 
-    def __init__(self, config: BertConfig, **kwargs):
+    def __init__(self, config: BertConfig, relations_definitions: List[Dict] = FEW_REL_DEFINITIONS, **kwargs):
         super().__init__(config)
         if not hasattr(config, "relation_emb_dim"):
             config.relation_emb_dim = 1024
         if not hasattr(config, "num_labels") or config.num_labels == 0:
             config.num_labels = 2
-        if not hasattr(config, "relations_definitions"):
-            self.relations_definitions = FEW_REL_DEFINITIONS
+        self.relations_definitions = relations_definitions
         self.bert = BertModel(config, **kwargs)
         self.num_labels = config.num_labels
         self.relation_emb_dim = config.relation_emb_dim