From 254d3d9da3446409309713abb62389d582825bbc Mon Sep 17 00:00:00 2001 From: Sho Hoshino Date: Mon, 26 Nov 2018 23:57:20 +0900 Subject: [PATCH] [WIP] Fix attention option when using fast beam search (#1073) Fix attention option when using fast beam search --- onmt/translate/translator.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/onmt/translate/translator.py b/onmt/translate/translator.py index 2166fa0822..5716cec374 100644 --- a/onmt/translate/translator.py +++ b/onmt/translate/translator.py @@ -188,7 +188,8 @@ def translate(self, all_predictions = [] for batch in data_iter: - batch_data = self.translate_batch(batch, data, fast=self.fast) + batch_data = self.translate_batch(batch, data, attn_debug, + fast=self.fast) translations = builder.from_batch(batch_data) for trans in translations: @@ -268,7 +269,7 @@ def translate(self, codecs.open(self.dump_beam, 'w', 'utf-8')) return all_scores, all_predictions - def translate_batch(self, batch, data, fast=False): + def translate_batch(self, batch, data, attn_debug, fast=False): """ Translate a batch of sentences. @@ -290,7 +291,7 @@ def translate_batch(self, batch, data, fast=False): self.max_length, min_length=self.min_length, n_best=self.n_best, - return_attention=self.replace_unk) + return_attention=attn_debug or self.replace_unk) else: return self._translate_batch(batch, data)