We will implement the skipgram model with pytorch.

Here is a great resource for understanding the skip gram model.

We will divide this post into three parts:

1. Loading and preparing dataset
2. Creating dataset tuples
3. Creating model
4. Training it

1. Loading and preparing dataset

For our task in creating word vectors we will use the movie plot description of wikipedia, available at https://www.kaggle.com/jrobischon/wikipedia-movie-plots . We will use the following code:

    from string import punctuation
    import pandas as pd
    df = pd.read_csv("data/wiki_movie_plots_deduped.csv")
    clear_punct_regex = "[" + punctuation + "\d\r\n]"
    corpus = df['Plot'].str.replace(clear_punct_regex, "").str.lower()
    corpus = " ".join(corpus)
    open("corpus2.txt", "w", encoding="utf8").write(corpus)

First we import pandas for parsing the csv file and them the punctuation variable that olds common punctuations used in strings.

>> punctuation
'!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~'

Line 4 we load the data of movie plot; line 5 we build the regex string used for remove punctuations from the text as well as control the characters: \n (new line), \r (also used for new line) and \d (any digits, 1 2 3 etc); line 6 we apply the regex string replacing every match with an empty string; line 7 we join all the rows of the movie plots strings into an unique long row; line 8 we write this cleaned long row of movie plots in a .txt file

We are done with data loading and cleaning.

2. Creating data tuples

    corpus = open("data/corpus.txt", encoding="utf8").readlines()
    corpus = " ".join(corpus).replace("\n", "")
    corpus = corpus.split(" ")

After creating the corpus file, we load it and remove line termination symbols from it (\n). We then split it into a tokens list.

    from collections import Counter
    vocab_cnt = Counter()
    vocab_cnt.update(corpus)
    vocab_cnt = Counter({w:c for w,c in vocab_cnt.items() if c > 5})

Next we count the number of ocurrences of each word and remove those that ocurrs less than 5 times.

    import numpy as np
    import random
    vocab = set()
    unigram_dist = list()
    word2id = dict()
    for i, (w, c) in enumerate(vocab_cnt.most_common()):
        vocab.add(w)
        unigram_dist.append(c)
        word2id[w] = i

    unigram_dist = np.array(unigram_dist)
    word_freq = unigram_dist / unigram_dist.sum()
    #Generate word frequencies to use with negative sampling
    w_freq_neg_samp = unigram_dist ** 0.75
    w_freq_neg_samp /= w_freq_neg_samp.sum() #normalize
    #Get words drop prob
    w_drop_p = 1 - np.sqrt(0.00001/word_freq)
    #Generate train corpus dropping common words
    train_corpus = [w for w in corpus if w in vocab and random.random() > w_drop_p[word2id[w]]]

In the above code we do several things:

• Create a vocab variable to stores all the words that appears in our corpus;
• Create a unigram_dist var to accumulate the number of occurrences of each word;
• Create a word2id dict that will help us encoding our words into values;
• Create w_freq_neg_samp, a probability distribution for sampling words in the negative sampling step;
• Create w_drop_p, a probability distribution for sampling words that will be in our training data;
• Create our training data, discarding some words based on w_drop_p.

    import torch
    #Generate dataset
    dataset = list()
    window_size = 5
    for i, w in enumerate(train_corpus):
        window_start = max(i - window_size, 0)
        window_end = i + window_size
        for c in train_corpus[window_start:window_end]:
            if c != w:
                dataset.append((word2id[w], word2id[c]))

    dataset = torch.LongTensor(dataset)
    if USE_CUDA:
        dataset = dataset.cuda()

In the above snippet we create the tuples of word and context word that will be used for training our model. We convert them in torch tensors and attach to gpus if they are available.

3. Creating the model

    import torch
    from torch import nn, optim
    import torch.nn.functional as F
    VOCAB_SIZE = len(word2id)
    EMBED_DIM = 128
    class Word2Vec(nn.Module):
        def __init__(self, vocabulary_size, embedding_dimension, sparse_grad=False):
            super(Word2Vec, self).__init__()

            self.embed_in = nn.Embedding(vocabulary_size, embedding_dimension, sparse=sparse_grad)
            self.embed_out = nn.Embedding(vocabulary_size, embedding_dimension, sparse=sparse_grad)
            #Sparse gradients do not work with momentum

            self.embed_in.weight.data.uniform_(-1, 1)
            self.embed_out.weight.data.uniform_(-1, 1)

        def neg_samp_loss(self, in_idx, pos_out_idx, neg_out_idxs):
            emb_in = self.embed_in(in_idx)
            emb_out = self.embed_out(pos_out_idx)

            pos_loss = torch.mul(emb_in, emb_out) #Perform dot product between the two embeddings by element-wise mult
            pos_loss = torch.sum(pos_loss, dim=1) #and sum the row values

            pos_loss = F.logsigmoid(pos_loss)

            neg_emb_out = self.embed_out(neg_out_idxs)
            #Here we must expand dimension for the input embedding in order to perform a matrix-matrix multiplication
            #with the negative embeddings
            neg_loss = torch.bmm(-neg_emb_out, emb_in.unsqueeze(2)).squeeze()
            neg_loss = F.logsigmoid(neg_loss)
            neg_loss = torch.sum(neg_loss, dim=1)

            total_loss = torch.mean(pos_loss + neg_loss)

            return -total_loss       

        def forward(self, indices):
            return self.embed_in(indices)

    w2v = Word2Vec(VOCAB_SIZE, EMBED_DIM, False)
    if USE_CUDA:
        w2v.cuda()

In the above class, we define our pytorch model. It is composed by two lookup embedding tables with uniform weight initialization.

To train our embeddings we are going to use Negative Sampling. In the function neg_samp_loss we compute the following quantity:

The first term

we compute between the lines 23 and 26; The second term

we compute between lines 28 and 33.

    def get_negative_samples(batch_size, n_samples):    
        neg_samples = np.random.choice(len(vocab), size=(batch_size, n_samples), replace=False, p=w_freq_neg_samp)

        if USE_CUDA:
            return torch.LongTensor(neg_samples).cuda()

        return torch.LongTensor(neg_samples)

Here we define our function to generate negative targets to be used in our objective during training.

    optimizer = optim.Adam(w2v.parameters(), lr=0.003)

Here we just define the optimizer that will perform our weight updates.

4. Training the model

    def get_batches(dataset, batch_size):
        for i in range(0, len(dataset), batch_size):
            yield dataset[i:i+batch_size]

This function is used for generate our batches during the training loop.

    n_epochs = 5
    n_neg_samples = 5
    batch_size = 512
    for epoch in range(n_epochs):  # loop over the dataset multiple times

        loss_values = []

        start_t = time.time()

        for dp in get_batches(dataset, batch_size):

            optimizer.zero_grad() # zero the parameter gradients

            inputs, labels = dp[:,0], dp[:,1]

            loss = w2v.neg_samp_loss(inputs, labels, get_negative_samples(len(inputs), n_neg_samples))
            loss.backward()

            optimizer.step()
            loss_values.append(loss.item())

        ellapsed_t = time.time() - start_t
        #if epoch % 1 == 0:
        print("{}/{}\tLoss: {}\tEllapsed time: {}".format(epoch + 1, n_epochs, np.mean(loss_values), ellapsed_t))
    print('Done')

5. Showing results

    import matplotlib.pyplot as plt
    from sklearn.manifold import TSNE
    data_viz_len = 300
    viz_embedding = w2v.embed_in.weight.data.cpu()[:data_viz_len]
    tsne = TSNE()
    embed_tsne = tsne.fit_transform(viz_embedding)
    plt.figure(figsize=(16,16))
    for w in vocab[:data_viz_len]:
        w_id = word2id[w]

        plt.scatter(embed_tsne[w_id,0], embed_tsne[w_id,1])
        plt.annotate(w, (embed_tsne[w_id,0], embed_tsne[w_id,1]), alpha=0.7)

If you have any questions, please ask them below, I will be happy to answer.