Minor modifications

yzhangcs · yzhangcs · commit 155704fee598 · 2020-12-13T15:17:07.000+08:00
diff --git a/supar/models/dependency.py b/supar/models/dependency.py
@@ -536,14 +536,78 @@ class CRF2oDependencyModel(BiaffineDependencyModel):
         https://www.aclweb.org/anthology/2020.acl-main.302/
     """
 
-    def __init__(self, n_lstm_hidden=400, n_mlp_sib=100, mlp_dropout=.33, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self,
+                 n_words,
+                 n_feats,
+                 n_rels,
+                 feat='char',
+                 n_embed=100,
+                 n_feat_embed=100,
+                 n_char_embed=50,
+                 bert=None,
+                 n_bert_layers=4,
+                 mix_dropout=.0,
+                 embed_dropout=.33,
+                 n_lstm_hidden=400,
+                 n_lstm_layers=3,
+                 lstm_dropout=.33,
+                 n_mlp_arc=500,
+                 n_mlp_sib=100,
+                 n_mlp_rel=100,
+                 mlp_dropout=.33,
+                 feat_pad_index=0,
+                 pad_index=0,
+                 unk_index=1,
+                 **kwargs):
+        super().__init__(**Config().update(locals()))
 
+        # the embedding layer
+        self.word_embed = nn.Embedding(num_embeddings=n_words,
+                                       embedding_dim=n_embed)
+        if feat == 'char':
+            self.feat_embed = CharLSTM(n_chars=n_feats,
+                                       n_embed=n_char_embed,
+                                       n_out=n_feat_embed,
+                                       pad_index=feat_pad_index)
+        elif feat == 'bert':
+            self.feat_embed = BertEmbedding(model=bert,
+                                            n_layers=n_bert_layers,
+                                            n_out=n_feat_embed,
+                                            pad_index=feat_pad_index,
+                                            dropout=mix_dropout)
+            self.n_feat_embed = self.feat_embed.n_out
+        elif feat == 'tag':
+            self.feat_embed = nn.Embedding(num_embeddings=n_feats,
+                                           embedding_dim=n_feat_embed)
+        else:
+            raise RuntimeError("The feat type should be in ['char', 'bert', 'tag'].")
+        self.embed_dropout = IndependentDropout(p=embed_dropout)
+
+        # the lstm layer
+        self.lstm = LSTM(input_size=n_embed+n_feat_embed,
+                         hidden_size=n_lstm_hidden,
+                         num_layers=n_lstm_layers,
+                         bidirectional=True,
+                         dropout=lstm_dropout)
+        self.lstm_dropout = SharedDropout(p=lstm_dropout)
+
+        # the MLP layers
+        self.mlp_arc_d = MLP(n_in=n_lstm_hidden*2, n_out=n_mlp_arc, dropout=mlp_dropout)
+        self.mlp_arc_h = MLP(n_in=n_lstm_hidden*2, n_out=n_mlp_arc, dropout=mlp_dropout)
         self.mlp_sib_s = MLP(n_in=n_lstm_hidden*2, n_out=n_mlp_sib, dropout=mlp_dropout)
         self.mlp_sib_d = MLP(n_in=n_lstm_hidden*2, n_out=n_mlp_sib, dropout=mlp_dropout)
         self.mlp_sib_h = MLP(n_in=n_lstm_hidden*2, n_out=n_mlp_sib, dropout=mlp_dropout)
+        self.mlp_rel_d = MLP(n_in=n_lstm_hidden*2, n_out=n_mlp_rel, dropout=mlp_dropout)
+        self.mlp_rel_h = MLP(n_in=n_lstm_hidden*2, n_out=n_mlp_rel, dropout=mlp_dropout)
 
+        # the Biaffine layers
+        self.arc_attn = Biaffine(n_in=n_mlp_arc, bias_x=True, bias_y=False)
         self.sib_attn = Triaffine(n_in=n_mlp_sib, bias_x=True, bias_y=True)
+        self.rel_attn = Biaffine(n_in=n_mlp_rel, n_out=n_rels, bias_x=True, bias_y=True)
+        self.criterion = nn.CrossEntropyLoss()
+        self.pad_index = pad_index
+        self.unk_index = unk_index
+
         self.crf = CRF2oDependency()
 
     def forward(self, words, feats):
diff --git a/supar/utils/config.py b/supar/utils/config.py
@@ -17,10 +17,10 @@ def __init__(self, conf=None, **kwargs):
                      **kwargs})
 
     def __repr__(self):
-        s = line = "-" * 15 + "-+-" + "-" * 25 + "\n"
-        s += f"{'Param':15} | {'Value':^25}\n" + line
+        s = line = "-" * 20 + "-+-" + "-" * 30 + "\n"
+        s += f"{'Param':20} | {'Value':^30}\n" + line
         for name, value in vars(self).items():
-            s += f"{name:15} | {str(value):^25}\n"
+            s += f"{name:20} | {str(value):^30}\n"
         s += line
 
         return s
diff --git a/supar/utils/field.py b/supar/utils/field.py
@@ -82,8 +82,7 @@ def __init__(self, name, pad=None, unk=None, bos=None, eos=None,
         self.tokenize = tokenize
         self.fn = fn
 
-        self.specials = [token for token in [pad, unk, bos, eos]
-                         if token is not None]
+        self.specials = [token for token in [pad, unk, bos, eos] if token is not None]
 
     def __repr__(self):
         s, params = f"({self.name}): {self.__class__.__name__}(", []
diff --git a/supar/utils/logging.py b/supar/utils/logging.py
@@ -32,7 +32,7 @@ def init_logger(logger,
 
 def progress_bar(iterator,
                  ncols=None,
-                 bar_format='{l_bar}{bar:36}| {n_fmt}/{total_fmt} {elapsed}<{remaining}, {rate_fmt}{postfix}',
+                 bar_format='{l_bar}{bar:18}| {n_fmt}/{total_fmt} {elapsed}<{remaining}, {rate_fmt}{postfix}',
                  leave=True):
     return tqdm(iterator,
                 ncols=ncols,
