Revise get_hitter_repr

microsoft · Nov 1, 2021 · b508106 · b508106
1 parent eedd02a
commit b508106
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/hitter-bert.py b/hitter-bert.py
@@ -230,7 +230,7 @@ def forward(self, batch):
         if self._hparams['use_hitter']:
             # kg_masks: [bs, 1, 1, length]
             # kg_embeds: nlayer*[bs, length, dim]
-            kg_embeds, kg_masks = self.hitter('get_hitter_repr', s=s, p=p)
+            kg_embeds, kg_masks = self.hitter('get_hitter_repr', s, p)
             kg_attentions = [None] * 2 + [(self.cross_attentions[i], kg_embeds[(i + 2) // 2], kg_masks)
                                           for i in range(self.kg_layer_num)]
         else:

diff --git a/kge/model/trme.py b/kge/model/trme.py
@@ -229,6 +229,9 @@ def forward(self, fn_name, *args, **kwargs):
         del self._scorer._entity_embedder
         del self._scorer._relation_embedder
 
+        if fn_name == 'get_hitter_repr':
+            return scores
+
         if self.training:
             self_loss_w = self.get_option("self_dropout")
             # MLM-like loss is weighted by the proportion of entities sampled