search

jarobyte91 / pytorch_beam_search

A lightweight implementation of Beam Search for sequence models in PyTorch.
https://pypi.org/project/pytorch-beam-search/
MIT License
49 stars 7 forks source link

Length Penalty with Beam Search #5

Open ra1995 opened 10 months ago

ra1995 commented 10 months ago

Hi, great work on the beam search implementation. I was going through your code to understand your implementation and noticed that you a) Don't terminate a beam when the EOS token is predicted by the language model, b) There is no length penalty for the probabilities.

Am I right in noticing these or have I missed something? I would be glad if you could give me an idea about how I can implement these two features. Thank you.

jarobyte91 commented 10 months ago

Hi,

You are right, there features are not implemented. I am aware that they are part of Beam Search, but I didn't need them for my work back then.

For a), I made it this way in order to keep the functions vectorized. For b), I couldn't find a way to implement it.

If you find a way to extend my code, I would be very happy to include them via a pull request.