Flexible aggregation for retrieval metrics

[daturkel]

🚀 Feature

Currently RetrievalMetric always computes the mean of the metric across all indexes (see this line). It would be nice to support other (or arbitrary) aggregation functions.

Motivation

I am using Retrieval metrics in a recommender model to aggregate per-user metrics. However, in one of our non-torch-based models, we're calculating per-user metrics and aggregating with median rather than mean. I'd like to be able to replicate this logic.

Pitch

RetrievalMetric can take a kwarg that specifies how to aggregate predictions. That kwarg could be a str from a fixed vocabulary of aggregation types (mean,median,min,max come to mind), or could allow the user to pass a callable that will take the predictions. I don't feel strongly about one approach vs the other. The API can default to mean to maintain backwards compatibility.

Alternatives

An alternative would be if the RetrievalMetric class maintains a buffer of computed predictions that I could access and aggregate myself. Currently, the compute method builds up a res list of metric results but doesn't save it anywhere—it just aggregates it and returns the aggregation. Exposing the full list of metric values would allow the user to a) aggregate them however they want, or b) inspect the distribution of the metric. This buffer would be cleared with a call to reset.

Additional context

If people like the idea, I could take a shot at implementing this.