Non-label features are not passed into data collator
I am trying to implement a customized data collator that tokenizes batches on-the-fly. But somehow,featurespassed to it does not have any field/key other thanlabel.  Hence, in my code below, when trying to extract the text for tokenization, I get the errorKeyError: 'text1'I read the source code of collators intransformersBut I don’t see where non-labelfields are dropped. Can someone help me?Below is my code which uses a hardcoded dataset of two keys/fields, namelytext1andlabel:from typing import List, Dict, Tuple, Literal, Any
import torch 
import datasets

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, DataCollatorWithPadding

train_ds = datasets.Dataset.from_dict(
                            {"text1":["A", "B", "C", "D", "E"], 
                             "label":[0, 0, 0, 1, 1]
                      })

class SmartCollator(DataCollatorWithPadding):
    """Tokenize each batch on the fly"""
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
    
    def __call__(self, features: List[Dict[str, Any]]):
        texts = [f["text1"] for f in features]
        labels = [f["label"] for f in features]

        encodings = self.tokenizer(texts, truncation=True, padding="longest", max_length=20, return_tensors="pt")

        return encodings.update({
            'labels': torch.tensor(labels)
        })

training_args = TrainingArguments(
    output_dir='./results',          
    num_train_epochs=1,              
    per_device_train_batch_size=2,
    report_to="none"
)

tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")

trainer = Trainer(
    model=model, 
    args=training_args,
    data_collator=SmartCollator(tokenizer=tokenizer),
    train_dataset=train_ds
)

trainer.train()
2025-01-01T02:32:32Z
There is probably a problem with the format of the data set. I think it will work as expected if you use thelist of dictionaries by Dataset.from_list().
2025-01-01T03:52:14Z
Thanks@John6666I triedfrom_list()Unfortunately, it does not work either.
2025-01-01T04:20:25Z
Oh… Perhaps this?stackoverflow.comQuestion about data_collator throwing a key error in Hugging facepython, pytorch, nlp, huggingface-transformersasked byChan Wingon05:42PM - 26 Sep 23 UTCStill error but improved maybe.from typing import List, Dict, Tuple, Literal, Any
import torch 
import datasets

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, DataCollatorWithPadding

#train_ds = datasets.Dataset.from_dict(
#                            {"text1":["A", "B", "C", "D", "E"], 
#                             "label":[0, 0, 0, 1, 1]
#                      })

train_ds = datasets.Dataset.from_list([{"text1": "A", "label": 0}, {"text1": "B", "label": 0}])

class SmartCollator(DataCollatorWithPadding):
    """Tokenize each batch on the fly"""
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
    
    def __call__(self, features: List[Dict[str, Any]]):
        print(features) # for debugging
        texts = [f["text1"] for f in features]
        labels = [f["label"] for f in features]

        encodings = self.tokenizer(texts, truncation=True, padding="longest", max_length=20, return_tensors="pt")

        return encodings.update({
            'labels': torch.tensor(labels)
        })

training_args = TrainingArguments(
    output_dir='./results',          
    num_train_epochs=1,              
    per_device_train_batch_size=2,
    report_to="none",
    remove_unused_columns=False,
)

tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=SmartCollator(tokenizer=tokenizer),
    train_dataset=train_ds
)

trainer.train()[{'text1': 'A', 'label': 0}, {'text1': 'B', 'label': 0}]
2025-01-01T07:05:55Z
I think I’ve solved it. dict.update() is a destructive method, so the return value is None. It cannot be returned.from typing import List, Dict, Tuple, Literal, Any
import torch 
import datasets

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, DataCollatorWithPadding

#train_ds = datasets.Dataset.from_dict(
#                            {"text1":["A", "B", "C", "D", "E"], 
#                             "label":[0, 0, 0, 1, 1]
#                      })

train_ds = datasets.Dataset.from_list([{"text1": "A", "label": 0}, {"text1": "B", "label": 0}])

class SmartCollator(DataCollatorWithPadding):
    """Tokenize each batch on the fly"""
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
    
    def __call__(self, features: List[Dict[str, Any]]):
        texts = [f["text1"] for f in features]
        labels = [f["label"] for f in features]

        encodings = self.tokenizer(texts, truncation=True, padding="longest", max_length=20, return_tensors="pt")
        encodings.update(labels=torch.tensor(labels))

        return encodings

training_args = TrainingArguments(
    output_dir='./results',          
    num_train_epochs=1,              
    per_device_train_batch_size=2,
    report_to="none",
    remove_unused_columns=False,
)

tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=SmartCollator(tokenizer=tokenizer),
    train_dataset=train_ds
)

trainer.train()
2025-01-01T18:28:44Z
You are right. I forgot that.update()for adictdoes not return the updated dict. Thanks so much.So to summarize:Besides theNoneTypereturn ofdict.update()in Python, the solution is that inTrainingArguments, setremove_unused_columns = Falseso columns unknown to theforward()method of the model (which I guess isinput_ids,attention_mask,labelsfor many models) will be preserved rather than removed by default.
2025-01-02T06:29:00Z
This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.