# Non-label features are not passed into data collator

**URL:** <https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791>\
**Category:** Beginners\
**Created:** [January 1, 2025, 12:27am UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791 "2025-01-01T00:27:00Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![forrestbao](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/forrestbao/32/23472_2.png) [@forrestbao](https://discuss.huggingface.co/u/forrestbao)\
**Post date:** [January 1, 2025, 12:27am UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791/1 "2025-01-01T00:27:00Z")

</div>

I am trying to implement a customized data collator that tokenizes batches on-the-fly. But somehow, `features` passed to it does not have any field/key other than `label`. Hence, in my code below, when trying to extract the text for tokenization, I get the error `KeyError: 'text1'`

I read the source code of collators in `transformers` But I don’t see where non-`label` fields are dropped. Can someone help me?

Below is my code which uses a hardcoded dataset of two keys/fields, namely `text1` and `label`:

```auto

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()

```

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [January 1, 2025, 2:32am UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791/2 "2025-01-01T02:32:32Z")

</div>

There is probably a problem with the format of the data set. I think it will work as expected if you use the **list of dictionaries by Dataset.from\_list()**.

---

<div class="post-metadata">

**Author:** ![forrestbao](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/forrestbao/32/23472_2.png) [@forrestbao](https://discuss.huggingface.co/u/forrestbao)\
**Post date:** [January 1, 2025, 3:52am UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791/3 "2025-01-01T03:52:14Z")

</div>

Thanks @John6666 I tried `from_list()` Unfortunately, it does not work either.

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [January 1, 2025, 4:20am UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791/4 "2025-01-01T04:20:25Z")

</div>

Oh… Perhaps this?

> <https://stackoverflow.com/questions/77182311/question-about-data-collator-throwing-a-key-error-in-hugging-face>

Still error but improved maybe.

```py
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()

```

```auto
[{'text1': 'A', 'label': 0}, {'text1': 'B', 'label': 0}]

```

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [January 1, 2025, 7:05am UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791/5 "2025-01-01T07:05:55Z")

</div>

I think I’ve solved it. dict.update() is a destructive method, so the return value is None. It cannot be returned.

```py
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()

```

---

<div class="post-metadata">

**Author:** ![forrestbao](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/forrestbao/32/23472_2.png) [@forrestbao](https://discuss.huggingface.co/u/forrestbao)\
**Post date:** [January 1, 2025, 6:28pm UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791/6 "2025-01-01T18:28:44Z")

</div>

You are right. I forgot that `.update()` for a `dict` does not return the updated dict. Thanks so much.

So to summarize:  
Besides the `NoneType` return of `dict.update()` in Python, the solution is that in `TrainingArguments`, set `remove_unused_columns = False` so columns unknown to the `forward()` method of the model (which I guess is `input_ids`, `attention_mask`, `labels` for many models) will be preserved rather than removed by default.

---

<div class="post-metadata">

**Author:** ![system](https://us1.discourse-cdn.com/hellohellohello/original/2X/d/de4155eb4aa4108ecb32a1389d7cc37ae69f88b7.png) [@system](https://discuss.huggingface.co/u/system)\
**Post date:** [January 2, 2025, 6:29am UTC](https://discuss.huggingface.co/t/non-label-features-are-not-passed-into-data-collator/133791/7 "2025-01-02T06:29:00Z")

</div>

Closed on January 2, 2025, 6:29am UTC

This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.