# Create custom data\_collator for Huggingface Trainer

**URL:** <https://discuss.huggingface.co/t/create-custom-data-collator-for-huggingface-trainer/16988>\
**Category:** 🤗Transformers\
**Created:** [April 19, 2022, 7:27pm UTC](https://discuss.huggingface.co/t/create-custom-data-collator-for-huggingface-trainer/16988 "2022-04-19T19:27:35Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![krishnagarg09](https://avatars.discourse-cdn.com/v4/letter/k/edb3f5/32.png) [@krishnagarg09](https://discuss.huggingface.co/u/krishnagarg09)\
**Post date:** [April 19, 2022, 7:27pm UTC](https://discuss.huggingface.co/t/create-custom-data-collator-for-huggingface-trainer/16988/1 "2022-04-19T19:27:35Z")

</div>

I need to create a custom `data_collator` for finetuning with Huggingface `Trainer API`.

HuggingFace offers `DataCollatorForWholeWordMask` for masking whole words within the sentences with a given probability.

```auto
model_ckpt = "vinai/bertweet-base"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt, normalization=True)
data_collator = DataCollatorForWholeWordMask(tokenizer=tokenizer, mlm_probability=args.mlm_prob)

```

I am using the collator as `Trainer` argument:

```auto
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset['train'],
        eval_dataset=None,
        data_collator=data_collator,
    )
    trainer.train()

```

But in my usecase, a sample input looks like: `<sent1>.<sent2>`. I want to mask the tokens only in `<sent2>` and not in `<sent1>`. How can I go about it? Any pointers on getting started with it are also welcome.

---

<div class="post-metadata">

**Author:** ![nbroad](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/nbroad/32/26465_2.png) [@nbroad](https://discuss.huggingface.co/u/nbroad)\
**Post date:** [July 22, 2022, 3:00pm UTC](https://discuss.huggingface.co/t/create-custom-data-collator-for-huggingface-trainer/16988/2 "2022-07-22T15:00:31Z")

</div>

Subclass DataCollatorForWholeWordMask and override the [torch\_mask\_tokens](https://github.com/huggingface/transformers/blob/d95a32cc60e5d92b4bf08cd805c6b0db7b4100cc/src/transformers/data/data_collator.py#L992-L1030) function.