EnhancementHelp WantedMediumnever-stale
Description
We have quite a few instances where we have some per-token losses/metrics along with a corresponding mask
metric_per_token # torch.Tensor of shape (batchsize, num_tokens)
mask # torch.BoolTensor of shape (batchsize, num_tokens)
And we want a per-batch-example average:
tokens_per_ex = mask.long().sum(dim=-1)
metric_per_ex = (metric_per_token * mask).sum(dim=-1)
metrics: List[MyMetric] = MyMetric.many(metric_per_ex, tokens_per_ex)
self.record_local_metric('metric_name', metrics)
I'd like us to have a helper classmethod in Metric called from_mask:
class Metric:
@classmethod
def from_mask(cls, metric_per_token, token_mask):
# returns the equivalent of the "metrics" object above
Once this is done, add unit tests for this (test AverageMetric and PPLMetric directly). Checkpoint there.
After you've implemented this, upgrade TorchGeneratorAgent to use your new helper, upgrading the code for loss, ppl, and token_acc.
See if you can find at least one other place who can benefit from upgrading this pattern.