jojonki / word2vec-pytorch

word2vec-pytorch

This repository shows an example of CBOW and Skip-gram (negative sampling version) known as Word2Vec algorithms.

word2vec-pytorch's Issues

Skipgram code is wrong, please refer to the original paper.

The original word2vec use two embeddings for each word, input embedding and output embedding. Your implementation only uses one.

how to extract the embedding vec for a given word?

Thank you for sharing this toy code. It is very informative to better understand the algorithm of word2vec. However, I encounter a problem of extract the word vectors for a given word. Usually, it could be done by adding a abstract function in the SkipGram class. Your suggestion would be appreciated a lot. Thanks.

class CBOW(nn.Module):
    def __init__(self, vocab_size, embd_size, context_size, hidden_size):
        super(CBOW, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embd_size)
        self.linear1 = nn.Linear(2*context_size*embd_size, hidden_size)
        self.linear2 = nn.Linear(hidden_size, vocab_size)
        
    def forward(self, inputs):
        embedded = self.embeddings(inputs).view((1, -1))
        hid = F.relu(self.linear1(embedded))
        out = self.linear2(hid)
        log_probs = F.log_softmax(out, dim = 1)
        return log_probs
    
    def extract(self, inputs):
        embeds = self.embeddings(inputs)
        return embeds


class SkipGram(nn.Module):
    def __init__(self, vocab_size, embd_size):
        super(SkipGram, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embd_size)
    
    def forward(self, focus, context):
        embed_focus = self.embeddings(focus).view((1, -1)) # input
        embed_ctx = self.embeddings(context).view((1, -1)) # output
        score = torch.mm(embed_focus, torch.t(embed_ctx)) # input*output
        log_probs = F.logsigmoid(score) # sigmoid
        return log_probs
    
    def extract(self, focus):
        embed_focus = self.embeddings(focus)
        return embed_focus