length pen param
@@ -302,6 +302,12 @@ class TorchGeneratorAgent(TorchAgent, ABC):
default=-1,
help='Size n-grams to block in beam search. val <= 0 implies no blocking',
)
+ agent.add_argument(
+ '--beam-length-penalty',
+ type=float,
+ default=0.65,
+ help='Applies a length penalty. Set to 0 for no penalty.',
+ )
agent.add_argument(
'--skip-generation',
type='bool',
@@ -804,6 +810,7 @@ class TorchGeneratorAgent(TorchAgent, ABC):
min_length=0,
block_ngram=self.beam_block_ngram,
context_block_ngram=self.beam_context_block_ngram,
+ length_penalty=self.opt.get('beam_length_penalty', 0.65),
padding_token=self.NULL_IDX,
bos_token=self.START_IDX,
eos_token=self.END_IDX,
@@ -815,6 +822,7 @@ class TorchGeneratorAgent(TorchAgent, ABC):
min_length=self.beam_min_length,
block_ngram=self.beam_block_ngram,
context_block_ngram=self.beam_context_block_ngram,
+ length_penalty=self.opt.get('beam_length_penalty', 0.65),
padding_token=self.NULL_IDX,
bos_token=self.START_IDX,
eos_token=self.END_IDX,
@@ -827,6 +835,7 @@ class TorchGeneratorAgent(TorchAgent, ABC):
min_length=self.beam_min_length,
block_ngram=self.beam_block_ngram,
context_block_ngram=self.beam_context_block_ngram,
+ length_penalty=self.opt.get('beam_length_penalty', 0.65),
padding_token=self.NULL_IDX,
bos_token=self.START_IDX,
eos_token=self.END_IDX,
@@ -839,6 +848,7 @@ class TorchGeneratorAgent(TorchAgent, ABC):
min_length=self.beam_min_length,
block_ngram=self.beam_block_ngram,
context_block_ngram=self.beam_context_block_ngram,
+ length_penalty=self.opt.get('beam_length_penalty', 0.65),
padding_token=self.NULL_IDX,
bos_token=self.START_IDX,
eos_token=self.END_IDX,
@@ -976,6 +986,7 @@ class TreeSearch(object):
eos_token=2,
min_length=3,
device='cpu',
+ length_penalty=0.65,
):
"""
Instantiate Beam object.
@@ -998,6 +1009,7 @@ class TreeSearch(object):
What device to use for computations
"""
self.beam_size = beam_size
+ self.length_penalty = length_penalty
self.block_ngram = block_ngram
self.min_length = min_length
self.eos = eos_token
@@ -1237,7 +1249,7 @@ class TreeSearch(object):
for finished_item in self.finished:
current_length = finished_item.timestep + 1
# these weights are from Google NMT paper
- length_penalty = math.pow((1 + current_length) / 6, 0.65)
+ length_penalty = math.pow((1 + current_length) / 6, self.length_penalty)
rescored_finished.append(
_HypothesisTail(
timestep=finished_item.timestep,
