Stabilize torch.topk() behavior

[melihyilmaz]

Addresses #284

To make sure we get the padding token at index '0' as the top scoring token for finished beams at each decoding step, I added a small epsilon of 1e-8 to index '0' in finished_mask so that it's not zeroed out like the rest of values on the same row. Again, these zero rows correspond to finished beams and we use masking to avoid extending them with new AA tokens.

This seems to resolve the error with minimal/no overhead (I get the same output on both CPU and GPU for the problematic mgf files mentioned in the issue) but unit tests still need to be added (@bittremieux feel free if you get a chance) and I'm open to moving away from torch.topk() if there are suggestions for a more robust or elegant solution.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (f01c607) 89.74% compared to head (7da1e2b) 89.09%.

❗ Current head 7da1e2b differs from pull request most recent head c58817e. Consider uploading reports for the commit c58817e to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##              dev     #290      +/-   ##
==========================================
- Coverage   89.74%   89.09%   -0.66%     
==========================================
  Files          12       12              
  Lines         917      917              
==========================================
- Hits          823      817       -6     
- Misses         94      100       +6     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[bittremieux]

I still don't find it super elegant, but it might be a practical fix for now. I combined some code to remove a few redundant steps, but the fix is the same.

@melihyilmaz Can you try to come up with a unit test that failed before the fix and runs now? So we can ensure that this is properly tested and avoid regression issues in the future.

[melihyilmaz]

I add a unit test for _get_topk_beams() that only passes with the current fix. I also compared results for a small mgf on CPU vs. GPU and those are identical now.

[lutfia95]

Hey, how can I run casanovo from source code? I had same problem under CPU, as the fix is in dev branch, I would like to use it before the new release.

[wsnoble]

You should be able to install from the dev branch like this:

pip install git+https://github.com/Noble-Lab/casanovo.git@dev

[bittremieux]

@lutfia95 We've now released Casanovo v4.1.0 as well that includes this fix, so you can more conveniently now upgrade from PyPI as well.