PreTrainedEncoderDecoder keeps giving me the same next token
Author: xiankgx
Pull request: no

Hi, I am trying to use PreTrainedEncoderDecoder to train a seq2seq model. I have a working training code, but I'm not sure if I'm doing things correctly because the decoded token is always the same token during inference.

The data is paired input and target sentence pairs.

The dataset class looks like this:
```
class LineByLineLabelledTextDataset(Dataset):
    """Labelled text dataset where a line corresponds to a sample."""

    def __init__(self,
                 lines,
                 tokenizer,
                 sep="|||",
                 max_seqlen=512):
        self.lines = lines
        self.tokenizer = tokenizer
        self.sep = sep
        self.max_seqlen = max_seqlen

    def __len__(self):
        return len(self.lines)

    def __getitem__(self, i):
        splitted = self.lines[i].split(self.sep)
        input, target = splitted[0], splitted[1]

        # target += " [GEN_STOP]"

        input_dict = self.tokenizer.encode_plus(input,
                                                max_length=self.max_seqlen,
                                                pad_to_max_length=True)
        target_dict = self.tokenizer.encode_plus(target,
                                                 max_length=self.max_seqlen,
                                                 pad_to_max_length=True)

        return torch.tensor(input_dict["input_ids"]), torch.tensor(target_dict["input_ids"]), torch.tensor(input_dict["attention_mask"]), torch.tensor(target_dict["attention_mask"])

```

The training function for one step looks like this:
```
def train_batch(batch, model, optimizer, device, phase="train"):
    input_ids = batch[0].to(device)
    target_ids = batch[1].to(device)
    input_attention_mask = batch[2].to(device)
    target_attention_mask = batch[3].to(device)

    optimizer.zero_grad()

    with torch.set_grad_enabled(phase == "train"):
        outputs = model(input_ids, target_ids,
                        encoder_attention_mask=input_attention_mask,
                        decoder_attention_mask=target_attention_mask,
                        decoder_lm_labels=target_ids)
        lm_loss = outputs[0]
        loss = lm_loss

    loss.backward()
    optimizer.step()

    return loss
```

The decode function looks like this
```
def decode(encoder_input_text, model, tokenizer, max_length=20):
    model.eval()
    
    text = encoder_input_text
    generated_text = "[CLS]"
    
    while len(generated_text.split()) < max_length:
        encoder_input_ids = tokenizer.encode(text)
        encoder_input_tensor = torch.tensor([encoder_input_ids])
        print(f"encoder_input_tensor: {encoder_input_tensor}")

        decoder_input_ids = tokenizer.encode(generated_text, add_special_tokens=False)
        decoder_input_tensor = torch.tensor([decoder_input_ids])
        print(f"decoder_input_tensor: {decoder_input_tensor}")
        
        
        with torch.no_grad():
            outputs = model(encoder_input_ids=encoder_input_tensor, decoder_input_ids=decoder_input_tensor)
            predictions = outputs[0]
            
            predicted_index = torch.argmax(predictions[0, -1]).item()
            predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])[0]
            
        generated_text += " " + predicted_token
            
        print(generated_text)
        print(len(generated_text.split()))
            
        if len(generated_text.split()) >= max_length:
            break
                
    return generated_text
```

I see the training loss goes down a bit during training. I don't know what I'm doing wrong.

Comments:
I'm using both "bert-base-uncased" for both encoder and decoder.

> I think I found the problem. I moved the code inside the train_step outside to the enclosing function and it seems to work.

Hi, I am having the same problem, what solved it for you?

> > I think I found the problem. I moved the code inside the train_step outside to the enclosing function and it seems to work.
> 
> Hi, I am having the same problem, what solved it for you?

Hi, the last time I tried, I was able to got it to work by training it at a lower learning rate and training for more iterations. Try troubleshooting the code by lowering the number of samples, and try to overfit to the training set by training it for more iterations. The loss should go down to at least less than 0.6. Proceed with the full dataset only when things work.