SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Generate: get the correct beam index on eos token (#18851)

Browse Source
This commit is contained in:
Joao Gante
2022-09-05 19:35:47 +01:00
committed by GitHub
parent c6d3daba54
commit d4dbd7ca59
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/generation_beam_search.py
View File

@@ -259,7 +259,7 @@ class BeamSearchScorer(BeamScorer):
continue
if beam_indices is not None:
beam_index = beam_indices[batch_beam_idx]
beam_index = beam_index + (next_index,)
beam_index = beam_index + (batch_beam_idx,)
else:
beam_index = None

2
tests/generation/test_generation_beam_search.py
View File

@@ -172,7 +172,7 @@ class BeamSearchTester:
input_ids[correct_idx].tolist(), beam_scorer._beam_hyps[batch_idx].beams[0][1].tolist()
)
self.parent.assertListEqual(
expected_beam_indices + [next_indices[batch_idx, 1].item()],
expected_beam_indices + [correct_idx],
torch.tensor(beam_scorer._beam_hyps[batch_idx].beams[0][2]).tolist(),
)