Simplify context layer in span ranking SRL

hankcs · hankcs · commit d994f8e35a88 · 2022-02-23T10:27:27.000-05:00
diff --git a/hanlp/components/srl/span_rank/span_ranking_srl_model.py b/hanlp/components/srl/span_rank/span_ranking_srl_model.py
@@ -473,7 +473,7 @@ def __init__(self, config, embed: torch.nn.Module, context_layer: torch.nn.Modul
         self.embed = embed
         # Initialize context layer
         self.context_layer = context_layer
-        context_layer_output_dim = context_layer.get_output_dim()
+        context_layer_output_dim = context_layer.get_output_dim() if context_layer else self.word_embedding_dim
         self.decoder = SpanRankingSRLDecoder(context_layer_output_dim, label_space_size, config)
 
     def forward(self,
@@ -484,9 +484,10 @@ def forward(self,
 
         context_embeddings = self.embed(batch)
         context_embeddings = F.dropout(context_embeddings, self.lexical_dropout, self.training)
-        contextualized_embeddings = self.context_layer(context_embeddings, masks)
+        if self.context_layer:
+            context_embeddings = self.context_layer(context_embeddings, masks)
 
-        return self.decoder.decode(contextualized_embeddings, sent_lengths, masks, gold_arg_starts, gold_arg_ends,
+        return self.decoder.decode(context_embeddings, sent_lengths, masks, gold_arg_starts, gold_arg_ends,
                                    gold_arg_labels, gold_predicates)
 
     @staticmethod
