# How to train TFBertForMaskedLM with TFTrainer

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-tfbertformaskedlm-with-tftrainer%2F1403%3E%5C
**Category:** Intermediate\
**Created:** [October 7, 2020, 12:06pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-tfbertformaskedlm-with-tftrainer%2F1403 "2020-10-07T12:06:52Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![fteufel](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/fteufel/32/435_2.png) [@fteufel](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Ffteufel)%5C
**Post date:** [October 7, 2020, 12:06pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-tfbertformaskedlm-with-tftrainer%2F1403%2F1 "2020-10-07T12:06:52Z")

</div>

I want to do language modeling finetuning of a pretrained Bert on a custom corpus. I want to use TPUs on google cloud, so I want to work with `TFTrainer` to avoid writing own code and don’t worry about its performance. I can’t find any info on how the masking is actually supposed to be performed here, when using `Trainer` in Pytorch it seems that `DataCollatorForLanguageModeling` is taking care of this.

This is how my setup looks like:  
tokenizer = BertTokenizer.from\_pretrained(args.tokenizer\_name)

```
    ds = load_dataset('text', data_files=[args.train_data_file])
    dataset = ds['train'].map(lambda examples: tokenizer(examples['text']), batched=True)
    dataset.set_format(type='tensorflow', columns=['text'])
    features = {x: dataset[x].to_tensor(default_value=0, shape=[None, tokenizer.max_len]) for x in ['input_ids', 'token_type_ids', 'attention_mask']}
    tfdataset = tf.data.Dataset.from_tensor_slices(features).batch(32)

    trainer = TFTrainer(model, training_args, tfdataset, optimizer = (tfa.optimizers.LAMB, None))

    trainer.train()
    trainer.save_model(args.save_path)

```

I’m using the datasets library to load a line-by-line txt file and convert it to a tensorflow dataset, this then goes straight to the trainer at the moment. Where and how in this process is the masking supposed to be added?

---

<div class="post-metadata">

**Author:** ![bengul](https://avatars.discourse-cdn.com/v4/letter/b/838e76/32.png) [@bengul](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fbengul)%5C
**Post date:** [February 23, 2022, 2:38am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-tfbertformaskedlm-with-tftrainer%2F1403%2F2 "2022-02-23T02:38:39Z")

</div>

Hi, Just wanted to know if you were able to train the model using this code snippet?