HuggingFace text summarization input data format issue
I’m trying to fine-tune a model to perform text summarization. I’m using AutoModelForSeq2SeqLM.from_pretrained(), so the following applies to several models (e.g. T5, ProphetNet, BART).
I’ve created a class called CustomDataset, which is a subclass of torch.utils.Dataset. That class contains one field: samples - a list of dictionaries that have encodings and labels keys. Each of the values in each of those dictionaries is a torch.Tensor. Here’s what an entry in samples looks like:
{'encoding': tensor([[21603, 10, 188, 563, 1]]), 'label': tensor([[ 1919, 22003, 22, 7, 1]])}
Here’s how I’m attempting to fine-tune the model using Trainer:
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

training_args = TrainingArguments(&quot;test_trainer&quot;)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=data,
)
trainer.train()

The error I’m getting gets thrown on line 63 in transformers\data\data_collator.py. Here’s that line of code:
label = first[&quot;label&quot;].item() if isinstance(first[&quot;label&quot;], torch.Tensor) else first[&quot;label&quot;]
Here’s the error message:
ValueError: only one element tensors can be converted to Python scalars
I understand why the error message specifically is being thrown - the first[&quot;label&quot;] tensor isn’t a one-element tensor, and hence item() can’t be called on it. That’s not why I’m asking this question, though.
I’m assuming that I’m not passing the data correctly, but it seems to me that Trainer should take care of input_ids and decoder_input_ids on its own. I’ve tried to set those manually (passing the encodings as input_ids and the labels as decoder_input_ids) and the model can successfully perform inference, but I haven’t managed to fine-tune it. Where am I making a mistake and how do I fix it?

Using the name label_ids instead of label fixes the specific problem. label should be used if the label is either an int, a float or a one-element torch.Tensor. For tensors with multiple elements, use label_ids. See data_collator.py, lines 62-71 for details:
if &quot;label&quot; in first and first[&quot;label&quot;] is not None:
    label = first[&quot;label&quot;].item() if isinstance(first[&quot;label&quot;], torch.Tensor) else first[&quot;label&quot;]
    dtype = torch.long if isinstance(label, int) else torch.float
    batch[&quot;labels&quot;] = torch.tensor([f[&quot;label&quot;] for f in features], dtype=dtype)
elif &quot;label_ids&quot; in first and first[&quot;label_ids&quot;] is not None:
    if isinstance(first[&quot;label_ids&quot;], torch.Tensor):
        batch[&quot;labels&quot;] = torch.stack([f[&quot;label_ids&quot;] for f in features])
    else:
        dtype = torch.long if type(first[&quot;label_ids&quot;][0]) is int else torch.float
        batch[&quot;labels&quot;] = torch.tensor([f[&quot;label_ids&quot;] for f in features], dtype=dtype)

Also, the name input_ids should be used instead of encoding. Otherwise, an unknown kwarg error gets thrown.
