Skip to content

Enable beam search#93

Merged
michaelbenayoun merged 6 commits intohuggingface:mainfrom
bocchris-aws:main
Jun 12, 2023
Merged

Enable beam search#93
michaelbenayoun merged 6 commits intohuggingface:mainfrom
bocchris-aws:main

Conversation

@bocchris-aws
Copy link
Contributor

Added beam search functionality with a test for bart-base model. Currently, KV cache does not work which is likely due to BART-specific preparation of decoder mask.

@bocchris-aws
Copy link
Contributor Author

@michaelbenayoun Would you mind taking a look?

@is_trainium_test
@pytest.mark.parametrize("model_name, use_cache, decoder_only", beam_search_testdata)
def test_beam_search_decoding(model_name, use_cache, decoder_only):
os.environ['NEURON_CC_FLAGS'] = '--model-type=transformer --enable-saturate-infinity'
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we need that, we should add this to the generate method no?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is actually specific to the BART model.

@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

@michaelbenayoun
Copy link
Member

You can fix the styling issue by running:

make style

@bocchris-aws
Copy link
Contributor Author

@michaelbenayoun is there anything on my end that I can do to make the tests pass?

@michaelbenayoun
Copy link
Member

No, I think this is ok, you would need to be a collaborator to be able to run the CI.

1 similar comment
@michaelbenayoun
Copy link
Member

No, I think this is ok, you would need to be a collaborator to be able to run the CI.

@michaelbenayoun michaelbenayoun merged commit 4f79926 into huggingface:main Jun 12, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants