deepmodeling / deepmodeling/unimol_tools
[Code scan] Handle no-padding pretraining batches in coordinate prediction
- Dominant language
- Python
- Stars
- 34
- Forks
- 9
- PR merge metrics
- No merged PRs in 30d
Description
This issue is a result of a Codex global repository scan.
## Summary
The pretraining model sets padding_mask to None when a batch has no padding. The coordinate prediction branch later unconditionally calls padding_mask.float(), so batch size 1 or equal-length batches can crash when masked_coord_loss is enabled.
## Code references
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/unimol.py#L71-L73
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/unimol.py#L100-L116
## Impact
Valid pretraining batches with no padded tokens can raise AttributeError before loss computation. This is especially likely for small datasets, homogeneous molecule sizes, or batch_size=1.
## Suggested fix
Keep a separate token-valid mask before nulling padding_mask for the encoder, or branch in the coordinate head and use an all-ones pair mask when padding_mask is None.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in unimol_tools/pretrain/unimol.py at lines 71-73 and 100-116, following how padding_mask is handled before the coordinate prediction loss. Verify the no-padding path with masked_coord_loss enabled, including batch_size=1 or equal-length batches. Done means valid batches without padding reach loss computation without an AttributeError while padded batches retain their existing behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 82/100