

TITLE
Make signature of `compute_metrics` parameter in Trainer class more flexible

COMMENTS
6

REACTIONS
+1: 1
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
# 🚀 Feature request

Current typing signature for the `compute_metrics` parameter in the `Trainer` class is:

```python
class Trainer:
    ...
    def __init__(
        ...
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        ...
```

As it is described now in with the Python typing system, the only parameter that you can pass to the function is `EvalPrediction`, containing the model predictions to calculate the metrics. I propose to make the function signature of `compute_metrics` a little bit more flexible, for example:

```python
        compute_metrics: Optional[Callable[[EvalPrediction, Optional[Any]], Dict]] = None,
```
or 
```python
        compute_metrics: Optional[Callable[[EvalPrediction, Optional[Dict]], Dict]] = None,
```

so users can pass an extra argument -e.g. a Dict- with additional information that can be used in the function.

Solution is not perfect in the sense that the typing system check of IDEs will scream when already defined functions in current user projects, pass a single parameter (see example below) as I haven't found a way of assigning a default value to the `Optional` in the `Callable` signature; using the `Ellipsis` is not possible either unless I've missed something (comments are welcome on this!!!)

## Motivation

For some models, I had to either pass some extra arguments to perform the metrics calculation or to access remote services to retrieve some additional data. The current typing signature for `compute_metrics` does not allow to pass these extra params so I had to do dirty workarounds.

```python
from abc import ABC
from typing import Optional, Callable, Dict, Any


def g(a:int):
    print(f"a in g: {a}")
    return {}


def h(a:int, b: Optional[int] = None):
    print(f"a in h: {a}")
    if b:
        print(f"b passed to g: {b}")
    return {}


class Dummy(ABC):

    def __init__(self, f: Optional[Callable[[int, Optional[Any]], Dict]] = None):
        self.f = f

    def test_f(self):
        if self.f:
            print(f"Calling {self.f}")
            if self.f.__name__ == "g":
                self.f(1)  # <- here the typing system screams a bit
            elif self.f.__name__ == "h":
                self.f(2,3)
            else:
                print("Not calling anything")


if __name__ == '__main__':
    o = Dummy(g)
    o.test_f()

    o = Dummy(h)
    o.test_f()

```
## Contribution

If someone else has had similar needs, you think this is a good idea, or you have better suggestion, I can provide a PR for this.