fix zero div and t5 decoder (#16)

castorini · May 2, 2020 · 3469584 · 3469584
1 parent 18f848a
commit 3469584
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pygaggle/model/evaluate.py b/pygaggle/model/evaluate.py
@@ -75,7 +75,7 @@ def accumulate(self, scores: List[float], gold: RelevanceExample):
         score_rels = self.truncated_rels(scores)
         score_rels[score_rels != 0] = 1
         gold_rels = np.array(gold.labels, dtype=int)
-        score = recall_score(gold_rels, score_rels, zero_division=1)
+        score = recall_score(gold_rels, score_rels, zero_division=0)
         self.scores.append(score)
 
 

diff --git a/pygaggle/rerank/transformer.py b/pygaggle/rerank/transformer.py
@@ -30,7 +30,7 @@ def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
             attn_mask = batch.output['attention_mask']
             _, batch_scores = greedy_decode(self.model,
                                             input_ids.to(self.device),
-                                            length=2,
+                                            length=1,
                                             attention_mask=attn_mask.to(self.device),
                                             return_last_logits=True)