Defining LLM Dataset types in Trainers or during Training Workflow

[mimipynb]

Feature request

Hi,
Just wondering if its possible to define a way to define the dataset input labels mapping argument in the transformers.Trainer or something similar?

I understand that the input data labels are dependent on the type of training loss function used and problem for example, the type of defining attention mask ids for Mask learner vs Causal learners etc. Would also be extremely helpful if I can have suggestions on how I can check what dataset column labels the trainer/model will read? -- besides the help(model.forward) as this really just shows the model's inputs as it is into the learner.

Any suggestions would be greatly appreciated! thanks!

Motivation

This problem is giving me clinical depression

Your contribution

If there's any suggestions in solving this case, I would love to contribute

[Rocketknight1]

cc @muellerzr @SunMarc

[SunMarc]

Is there a specific issue or bug that you are tracking in your code @mimipynb ?