ipsumdominum / pytorch-simple-transformer Goto Github PK

###################
## original leave-last-token out decoder, 
## Not sure what's the exact error in this calculation, 
## but maybe because the model see the mask token directly?
##
#Output German, One Token At A Time
all_outs = torch.tensor([],requires_grad=True).to(device)
for i in range(item["german"].shape[1]-1):
    out = model(item["german"][:,:i+1])
    all_outs = torch.cat((all_outs,out),dim=1)


# ###################
# My variation of leave-last-token-out decoder, Used at training
# output_vocab_size = german_vocab_len

g = item["german"].shape
x = torch.zeros( [g[0],g[1],],dtype=torch.long ).to(device)
all_outs = torch.tensor([],requires_grad=True).to(device)
for i in range(item["german"].shape[1]-1):
    xx = torch.zeros( [g[0],g[1], ],dtype=torch.long ).to(device)
    out = model(x)
    xx[:,i:i+1] = item["german"][:,i:i+1]
    x = x+xx
    all_outs = torch.cat((all_outs,out),dim=1)

# ###################
# My variation of beam search decoder
model.encode(item["english"][:,1:-1])
g = item["german"].shape
x = torch.zeros( [g[0],g[1],],dtype=torch.long ).to(device)
all_outs = torch.tensor([],requires_grad=True).to(device)
for i in range(item["german"].shape[1]-1):
    out = model(x)
    x[:,i:i+1] = out.argmax(axis=-1)
    all_outs = torch.cat((all_outs,out),dim=1)

attention_weights = F.softmax(attention_weights,dim=2)
attention_weights = attention_weights *0.   ## Try this!

class TransformerTranslator(nn.Module):
    def __init__(self,embed_dim,num_blocks,num_heads,vocab_size,CUDA=False):
        super(TransformerTranslator,self).__init__()
        self.embedding = Embeddings(vocab_size,embed_dim,CUDA=CUDA)
        self.encoder = Encoder(embed_dim,num_heads,num_blocks,CUDA=CUDA)
        self.decoder = Decoder(embed_dim,num_heads,num_blocks,vocab_size,CUDA=CUDA)
        self.encoded = False
        self.device = torch.device('cuda:0' if CUDA else 'cpu')
    def encode(self,input_sequence):
        embedding = self.embedding(input_sequence).to(self.device)
        self.encode_out = self.encoder(embedding)
        self.encoded = True
    def forward(self,output_sequence):
        if(self.encoded==False):
            print("ERROR::TransformerTranslator:: MUST ENCODE FIRST.")
            return output_sequence
        else:
            embedding = self.embedding(output_sequence)
            return self.decoder(self.encode_out,embedding)