Pretraining for encoder of TF T5 model
Author: dharakotecha
Pull request: no
Labels: name: wontfix

Hi, for sequence to sequence task in tensorflow, I see that I can use TFT5 model. However, before sequence to sequence training, I need to perform masked language model pretraining on the encoder and initialise the weights of the encoder and decoder with the same weights that I obtain by masked language model pretraining. Is there a way to do this in the tensorflow version of T5?
I could do that with the EncoderDecoder, but it is only supported in PyTorch and not tensorflow.

Comments:
@dharakotecha, you would need to create a MaskedLM head similar to distilbert. You would also need to either use one of the masekd token <extra_id_0> or add a new token for mask. 

Some experimentation on the t5 encoder shows that it does produce pretty reasonable output, but does not do MLM at the encoder side. The cloze/fill in the blank happens in the decoder side.

This is just quick and dirty. You would need to create a module similar to a MLM for distilbert with the loss, feeding in the output from the encoder into the projector:

```
t5 = AutoModel.from_pretrained('t5-small').cuda()
tokenizer = AutoTokenizer.from_pretrained('t5-small')
encoder = t5.encoder
embeddings = encoder.embed_tokens
config = t5.config
#tied the embedding to the projector. 
vocab_projector = nn.Linear(config.d_model, config.vocab_size).cuda()
vocab_projector.weight = embeddings.weightvocab_projector = nn.Linear(config.d_model, config.vocab_size).cuda()
vocab_projector.bias.data = torch.nn.functional.pad(vocab_projector.bias.data, (0, vocab_projector.weight.shape[0] - vocab_projector.bias.shape[0],), "constant", 0)

def pred(predictions):
  for pred in predictions:
    print ("**")
    sorted_preds, sorted_idx = pred.sort(dim=-1, descending=True)
    ret = []
    for k in range(2):
        predicted_index = [sorted_idx[i, k].item() for i in range(0,len(predictions[0]))]
        predicted_token = ' '.join([tokenizer.convert_ids_to_tokens([predicted_index[x]])[0] for x in range(1,len(predictions[0]))]).replace('Ġ', ' ').replace('  ', ' ').replace('##', '')
        ret.append(predicted_token)
    return ret


input_txt = ["
Lincoln was an American president and lawyer"]
inputs = tokenizer(input_txt, return_tensors='pt', add_special_tokens=True, padding=True)
predictions = vocab_projector(encoder(inputs.input_ids.cuda())[0])
pred(predictions)


```
Outputs:
['\u2581Lincoln \u2581was \u2581psiho American \u2581President & \u2581lawyer 
', '\u2581senzati  tais Mitglied \u2581American \u2581president \u2581and \u2581Lawyer \u2581summarize']**

But, using the extra mask, will not infer the missing token.

```
input_txt = ["
Lincoln <extra_id_0> American president and lawyer"]
inputs = tokenizer(input_txt, return_tensors='pt', add_special_tokens=True, padding=True)
predictions = vocab_projector(encoder(inputs.input_ids.cuda())[0])
pred(predictions)

```
Wil output:
['\u2581Lincoln <extra_id_0> \u2581American \u2581president & \u2581lawyer 
', '\u2581Abraham \u2581botez American \u2581President \u2581and \u2581Lawyer gasesc']**

Hope this helps,



I just realized that you asking about tensor flow too... I'm assuming you could do the similar thing but creating a MLM head in tesnor flow just for the encoder.

This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.