Skip to content
This repository has been archived by the owner on Dec 29, 2022. It is now read-only.

#174 fixed that unk_replacement and beam_search work together #271

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

#174 fixed that unk_replacement and beam_search work together #271

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
7 changes: 6 additions & 1 deletion seq2seq/tasks/decode_text.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,8 @@ def before_run(self, _run_context):

if "attention_scores" in self._predictions:
fetches["attention_scores"] = self._predictions["attention_scores"]
elif "beam_search_output.original_outputs.attention_scores" in self._predictions:
fetches["beam_search_output.original_outputs.attention_scores"] = self._predictions["beam_search_output.original_outputs.attention_scores"]

return tf.train.SessionRunArgs(fetches)

Expand All @@ -169,7 +171,10 @@ def after_run(self, _run_context, run_values):
if self._unk_replace_fn is not None:
# We slice the attention scores so that we do not
# accidentially replace UNK with a SEQUENCE_END token
attention_scores = fetches["attention_scores"]
if "beam_search_output.original_outputs.attention_scores" in fetches:
attention_scores = fetches["beam_search_output.original_outputs.attention_scores"][:,0,:]
else:
attention_scores = fetches["attention_scores"]
attention_scores = attention_scores[:, :source_len - 1]
predicted_tokens = self._unk_replace_fn(
source_tokens=source_tokens,
Expand Down