How to write a generation function for text pytorch transformer?

Viewed 111

Following this pytorch tutorial, I'm able to create and train a transformer model on a custom dataset. The problem is, I've scoured the web and have found no clear answers... How do I use this model to generate text? I took a stab at it, by encoding my SOS and seed text and passing it through the model's forward method... But this only produces repeating garbage. The src_mask doesn't appear to be the right size or functioning at all.

def generate(model: nn.Module, src_text:str):
    src=BeatleSet.encode(src_text.lower()) # encodes seed text
    SOS=BeatleSet.textTokDict['<sos>'] ; EOS=BeatleSet.textTokDict['<eos>'] # obtain eos and sos tokens
    model.eval(); entry=[SOS]+src
    y_input=torch.tensor([entry], dtype=torch.long, device=device) # convert entry to tensor
    num_tokens=len(BeatleSet)
    for i in range(50):
        src_mask=generate_square_subsequent_mask(y_input.size(0)).to(device) #create a mask of size 1,1 (???)
        pred=model(y_input, src_mask) # passed through forward method
        next_item = pred.topk(1)[1].view(-1)[-1].item() # selecting the most probable next token (I think)
        next_item = torch.tensor([[next_item]], device=device)
        y_input=torch.cat((y_input, next_item), dim=1) # added to inputs to be run again
        if next_item.view(-1).item() == EOS:
            break
    return " ".join(BeatleSet.decode(y_input.view(-1).tolist()))
    
print(generate(model, "Oh yeah I"))

For the record, I'm following the architecture to the letter. This should be reproducible with the wikidata set that is used in the tutorial. Please advise, I've been banging my head on this one.

0 Answers
Related