Added gaussian label smoother for number tokens #34
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…se. Added more extensive testing of the gaussian smoother.
Merged
* Refractored number token loss so get a number token selector * Added first implementation of gaussian label smoother * Added test cases for the label smoother * Small change to args * added changes to number_token_loss and label_smoother * added number_token_selector.py * added init * changed trainer * Added comments, and valid mask before one-hot encoding * added a case differenciation for sigma=0 * For label_smoother: Fixed gradient flow for no valid number tokens case. Added more extensive testing of the gaussian smoother. * Fixed bug in test file * nvocab fix * Bigger commit, will probably fail the gce tests * Cleaned up * Cleaned up --------- Co-authored-by: ad045 <[email protected]>
jannisborn
approved these changes
Jan 7, 2025
There was a problem hiding this comment.
Great work @ad045 ! Seems very flexible and easy to use. Super solid job on the test suite 👍🏼
I left some cosmetics comments where unnecessary things can be removed, but conceptually we are ready to merge already now!
src/ntl/run_language_modeling.py
Outdated
| @@ -349,6 +364,8 @@ def run_language_modeling(model_args: ModelArguments, training_args: TrainingArg | |||
| train_dataset=train_dataset, | |||
| eval_dataset=eval_dataset, | |||
| tokenizer=tokenizer, | |||
| # selector=selector, | |||
src/ntl/run_language_modeling.py
Outdated
| selector=selector | ||
| ) | ||
| else: | ||
| selector = None |
There was a problem hiding this comment.
Remove selector since it is unused downstream
| self.nvocab[id] = self.tokenizer.decode_number_token(token, ignore_order=True) | ||
|
|
||
|
|
||
| def select_number_tokens(self, logits: Tensor, labels: Tensor): |
There was a problem hiding this comment.
Labels are not used in here so I would remove them from the signature
There was a problem hiding this comment.
and remove them from the return statement below then obviously
src/ntl/utils/label_smoother.py
Outdated
| raise AttributeError("The selector must have an attribute 'nvocab' representing the number of valid vocab tokens.") | ||
|
|
||
| # Select number tokens | ||
| logits, labels, number_tokens = self.selector.select_number_tokens(logits, labels) |
| # Create a mask to filter out non-digit tokens | ||
| number_tokens = ~torch.isnan(self.nvocab) | ||
| logits = logits[:, :, number_tokens] | ||
| logits, labels, number_tokens = self.selector.select_number_tokens(logits, labels) |
There was a problem hiding this comment.
remove labels here also
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description:
This pull request introduces a label smoothing technique that applies Gaussian smoothing exclusively to number tokens. The key changes are outlined below:
Gaussian Label Smoother:
GaussianLabelSmootherclass insrc/ntl/utils/label_smoother.py.ModelArguments:gaussian_label_smootherandlabel_smoother_sigma.Number Token Selector:
NumberTokenLossinto a newNumberTokenSelectorclass located insrc/ntl/utils/number_token_selector.py.GaussianLabelSmootherandNumberTokenLossto consistently select number tokens.Other Changes:
src/ntl/run_language_modeling.pyandsrc/ntl/trainer.pyto support the new label smoothing feature.src/ntl/loss_functions/number_token_loss.pyto integrate with the newNumberTokenSelector.src/ntl/args.pyto include new arguments for the Gaussian Label Smoother.GaussianLabelSmootherandNumberTokenSelector.Types of Changes