add estimator in greedy inference (#135)

eole-nlp · Oct 25, 2024 · 5bf241b · 5bf241b
1 parent d7959ba
commit 5bf241b
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 1 deletion.
diff --git a/eole/predict/encoder.py b/eole/predict/encoder.py
@@ -44,6 +44,7 @@ def predict_batch(self, batch, attn_debug):
                     top_p=self.top_p,
                     beam_size=self.beam_size,
                     ban_unk_token=self.ban_unk_token,
+                    add_estimator=self.add_estimator,
                 )
             else:
                 # TODO: support these blacklisted features
@@ -66,6 +67,7 @@ def predict_batch(self, batch, attn_debug):
                     stepwise_penalty=self.stepwise_penalty,
                     ratio=self.ratio,
                     ban_unk_token=self.ban_unk_token,
+                    add_estimator=self.add_estimator,
                 )
             return self._predict_batch_with_strategy(batch, decode_strategy)
 

diff --git a/eole/predict/generator.py b/eole/predict/generator.py
@@ -47,6 +47,7 @@ def predict_batch(self, batch, attn_debug, scoring=False):
                     top_p=self.top_p,
                     beam_size=self.beam_size,
                     ban_unk_token=self.ban_unk_token,
+                    add_estimator=self.add_estimator,
                 )
             else:
                 # TODO: support these blacklisted features

diff --git a/eole/predict/greedy_search.py b/eole/predict/greedy_search.py
@@ -162,6 +162,10 @@ def __init__(
         self.topk_scores = None
         self.beam_size = beam_size
         self.n_best = n_best
+        if add_estimator:
+            self.num_hyp = self.beam_size
+        else:
+            self.num_hyp = self.n_best
 
     def initialize(self, enc_out, src_len, device=None, target_prefix=None):
         """Initialize for decoding."""
@@ -282,7 +286,7 @@ def update_finished(self):
         if self.done:
             for b in range(self.batch_size):
                 best_hyp = sorted(self.hypotheses[b], key=lambda x: x[0], reverse=True)[
-                    : self.n_best
+                    : self.num_hyp
                 ]
                 for score, pred, attn in best_hyp:
                     self.scores[b].append(score)

diff --git a/eole/predict/translator.py b/eole/predict/translator.py
@@ -110,6 +110,7 @@ def predict_batch(self, batch, attn_debug):
                     top_p=self.top_p,
                     beam_size=self.beam_size,
                     ban_unk_token=self.ban_unk_token,
+                    add_estimator=self.add_estimator,
                 )
             else:
                 # TODO: support these blacklisted features