Support nested data structures for input data
Author: alex2awesome
Pull request: no
Labels: name: wontfix

# 🚀 Feature request

Support in `train.py` for more complicated/nested data inputs rather than the current assumed single-layer dictionary structure.

## Motivation

In many applications -- for example, mine, where I wish to implement a semi-supervised learning protocol -- the training input might need to be different than the 1-layer dictionary that `train.py` is currently hard-coded to accept.

For instance, in my case, I need a nested structure that supports:

```
    {'supervised_data': {'input_ids':[....], 'labels': [...], 'attention_mask'},
      'unsupervised_data': {'input_ids':[...], 'attention_mask'},
      'augmented_data': {'input_ids':[...], 'attention_mask'}
```
(I am attempting to implement the following paper, Unsupervised Data Augmentation for Consistency Training https://arxiv.org/abs/1904.12848).

However, I can imagine other use-cases including MAML, multi-task learning and multi-modal learning where huggingface would provide a great framework but is currently limited in it's data input format.

## Your contribution

I've identified a couple of quick fixes for this:

Lines 962-974 of `train.py`, or the `_prepare_inputs` function should be rewritten as:

```
    def _prepare_inputs(self, inputs):
        """
        Prepare :obj:`inputs` before feeding them to the model, converting them to tensors if they are not already and
        handling potential state.
        """
        def map_nested_dicts_modify(ob, func):
            if isinstance(ob, dict):
                return {k: map_nested_dicts_modify(v, func) for k, v in ob.items()}
            if isinstance(ob, list):
                return list(map(lambda x: map_nested_dicts_modify(x, func), ob))
            else:
                return func(ob)

        def to_device(v):
            if isinstance(v, torch.Tensor):
                v = v.to(self.args.device)
            return v

        inputs = map_nested_dicts_modify(inputs, to_device)

        if self.args.past_index >= 0 and self._past is not None:
            inputs["mems"] = self._past

        return inputs
```

Line 1135 of `train.py` should be expanded to:

```        
        def _finditem(obj, key):
            if key in obj: return True
            for k, v in obj.items():
                if isinstance(v, dict):
                    return _finditem(v, key)

        has_labels = any(_finditem(inputs, k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])
```

I'm not quite sure how to handle line 1254, or whether it is really that necessary, but one way might be to again use `_finditem` for `["labels", "input_ids"]`.

And that's it -- it's then up to the user to modify their models so that their `DataCollator` generates the expected structure, and `forward` takes in the expected structure!

Even if you explicitly don't wish to support these different training protocols mentioned above, I do think, from a software engineering perspective, that `train.py` should be more fully abstracted from particulars of the data inputs and model inputs. This feature takes you a step closer to that (although not completely, as the `has_labels` line still expects a certain set of keys _somewhere_ in the data input.) 

Better yet than the suggestions here would be to make the data input it's own class for full abstraction, but I can see the argument against you doing that, as it is yet another data class for users to learn and code up, and would be a breaking update for all those who have implemented `DataCollators` that adhere to your guidelines.

Alex

Comments:
Note that there are multiple frameworks that provide generic training loops. The goal of `Trainer` (I'm assuming you're talking about it since there is no `train.py` file) is not to replace them or compete with them but to provide an easy way to train and finetune Transformers models. Those models don't take nested inputs, so Trainer does not support this. Those models are expected to return the loss as the first item of their output, so Trainer expects it too.

Making Trainer more easily customizable by providing better hooks for subclassing (your use case could be done by overriding the two private methods you mention for instance) is something we are working on, but we won't have a base Trainer that is too generic, it will remain customized to the models the library provides.

Thank you for your consideration and comments! 

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.