Make signature of `compute_metrics` parameter in Trainer class more flexible
Author: francisco-perez-sorrosal
Pull request: no
Labels: name: wontfix

# 🚀 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.

Comments:
> 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.

How did you do that in the current `Trainer`? I'm not against making `compute_metrics` more flexible but I don't see how you can add more arguments to its call without subclassing Trainer and overriding certain methods, in which case, you can also override the init.

In fact I have my own subclassed Trainer, but I wanted to reuse as many as possible of the current `__init__` parameters of the regular Trainer. I know I can define my own `compute_metrics` function with a different signature in the `__init__` method of my Trainer, but I was trying to avoid that and reuse the current `compute_metrics` signature  :-)

Maybe the example above is not clear enough as the the DummyTrainer is not subclassing the standard one. The example was trying to highlight the function signature more than the reuse of the original Trainer.


This seems like a very edge case to me, so I would leave the current `Trainer` as is, and adapt the `__init__` in your subclasses of `Trainer`.

Yes, I agree that is kind of a corner case. In my subclassed Trainer I was trying to reuse the `prediction_loop()` method from the base Trainer but I needed to pass some more parameters to my `compute_metrics` function, apart from the `EvalPrediction` param. So the easiest workaround was of course, to copy the original `prediction_loop()` in my subclassed Trainer and, instead of calling:

https://github.com/huggingface/transformers/blob/8062fa63c564d4cc0d29573acd89092f1eb1df64/src/transformers/trainer.py#L1398

, I call my version of `compute_metrics` with the extra parameters. 

With this proposal was trying to avoid that ugly copy-paste I had to do by 1) changing the `compute_metrics` signature as I described above, 2) defining an attribute in Trainer to serve as a placeholder for the extra metrics arguments (and which can be set from subclassed Trainers,) e.g.:

```python
self.metrics_extra_args: Optional[Dict] = None
```
and 3) changing line 1398 to something like:

```python
metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids),  self.metrics_extra_args)
```

But I understand that is kind of convoluted. Thanks @sgugger  in any case for your consideration!

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

I really need this