Learning Word2Vec Subword Representations using BlazingText

Word2Vec is a popular algorithm used for generating dense vector representations of words in large corpora using unsupervised learning. These representations are useful for many natural language processing (NLP) tasks like sentiment analysis, named entity recognition and machine translation.

Popular models that learn such representations ignore the morphology of words, by assigning a distinct vector to each word. This is a limitation, especially for languages with large vocabularies and many rare words. SageMaker BlazingText can learn vector representations associated with character n-grams; representing words as the sum of these character n-grams representations [1]. This method enables BlazingText to generate vectors for out-of-vocabulary (OOV) words, as demonstrated in this notebook.

Popular tools like FastText learn subword embeddings to generate OOV word representations, but scale poorly as they can run only on CPUs. BlazingText extends the FastText model to leverage GPUs, thus providing more than 10x speedup, depending on the hardware.

[1] P. Bojanowski, E. Grave, A. Joulin, T. Mikolov, Enriching Word Vectors with Subword Information

Setup

Let’s start by specifying:

• The S3 bucket and prefix that you want to use for training and model data. This should be within the same region as the Notebook Instance, training, and hosting. If you don’t specify a bucket, SageMaker SDK will create a default bucket following a pre-defined naming convention in the same region.

• The IAM role ARN used to give SageMaker access to your data. It can be fetched using the get_execution_role method from sagemaker python SDK.

import sagemaker
from sagemaker import get_execution_role
import boto3
import json

sess = sagemaker.Session()

role = get_execution_role()
print(
    role
)  # This is the role that SageMaker would use to leverage AWS resources (S3, CloudWatch) on your behalf

bucket = sess.default_bucket()  # Replace with your own bucket name if needed
print(bucket)
prefix = "blazingtext/subwords"  # Replace with the prefix under which you want to store the data if needed

Data Ingestion

Next, we download a dataset from the web on which we want to train the word vectors. BlazingText expects a single preprocessed text file with space separated tokens and each line of the file should contain a single sentence.

In this example, let us train the vectors on text8 dataset (100 MB), which is a small (already preprocessed) version of Wikipedia dump.

s3 = boto3.client("s3")
s3.download_file("sagemaker-sample-files", "datasets/text/text8/text8.gz", "text8.gz")

# Uncompressing
!gzip -d text8.gz -f

After the data downloading and uncompressing is complete, we need to upload it to S3 so that it can be consumed by SageMaker to execute training jobs. We’ll use Python SDK to upload these two files to the bucket and prefix location that we have set above.

train_channel = prefix + "/train"

sess.upload_data(path="text8", bucket=bucket, key_prefix=train_channel)

s3_train_data = "s3://{}/{}".format(bucket, train_channel)

Next we need to setup an output location at S3, where the model artifact will be dumped. These artifacts are also the output of the algorithm’s training job.

s3_output_location = "s3://{}/{}/output".format(bucket, prefix)

Training Setup

Now that we are done with all the setup that is needed, we are ready to train our object detector. To begin, let us create a sageMaker.estimator.Estimator object. This estimator will launch the training job.

region_name = boto3.Session().region_name

container = sagemaker.image_uris.retrieve(
    region=region_name, framework="blazingtext", version="latest"
)
print("Using SageMaker BlazingText container: {} ({})".format(container, region_name))

Training the BlazingText model for generating word vectors

Similar to the original implementation of Word2Vec, SageMaker BlazingText provides an efficient implementation of the continuous bag-of-words (CBOW) and skip-gram architectures using Negative Sampling, on CPUs and additionally on GPU[s]. The GPU implementation uses highly optimized CUDA kernels. To learn more, please refer to *BlazingText: Scaling and Accelerating Word2Vec using Multiple GPUs*.

Besides skip-gram and CBOW, SageMaker BlazingText also supports the “Batch Skipgram” mode, which uses efficient mini-batching and matrix-matrix operations (BLAS Level 3 routines). This mode enables distributed word2vec training across multiple CPU nodes, allowing almost linear scale up of word2vec computation to process hundreds of millions of words per second. Please refer to *Parallelizing Word2Vec in Shared and Distributed Memory* to learn more.

BlazingText also supports a supervised mode for text classification. It extends the FastText text classifier to leverage GPU acceleration using custom CUDA kernels. The model can be trained on more than a billion words in a couple of minutes using a multi-core CPU or a GPU, while achieving performance on par with the state-of-the-art deep learning text classification algorithms. For more information, please refer to algorithm documentation or the text classification notebook.

To summarize, the following modes are supported by BlazingText on different types instances:

Modes	cbow (supports subwords training)	skipgram (supports subwords training)	batch_skipgram	supervised
Single CPU instance	✔	✔	✔	✔
Single GPU instance	✔	✔		✔ (Instance with 1 GPU only)
Multiple CPU instances			✔

Now, let’s define the resource configuration and hyperparameters to train word vectors on text8 dataset, using “skipgram” mode on a c4.2xlarge instance.

bt_model = sagemaker.estimator.Estimator(
    container,
    role,
    instance_count=1,
    instance_type="ml.c4.2xlarge",  # Use of ml.p3.2xlarge is highly recommended for highest speed and cost efficiency
    volume_size=30,
    max_run=360000,
    input_mode="File",
    output_path=s3_output_location,
    sagemaker_session=sess,
)

Please refer to algorithm documentation for the complete list of hyperparameters.

bt_model.set_hyperparameters(
    mode="skipgram",
    epochs=5,
    min_count=5,
    sampling_threshold=0.0001,
    learning_rate=0.05,
    window_size=5,
    vector_dim=100,
    negative_samples=5,
    subwords=True,  # Enables learning of subword embeddings for OOV word vector generation
    min_char=3,  # min length of char ngrams
    max_char=6,  # max length of char ngrams
    batch_size=11,  #  = (2*window_size + 1) (Preferred. Used only if mode is batch_skipgram)
    evaluation=True,
)  # Perform similarity evaluation on WS-353 dataset at the end of training

Now that the hyper-parameters are setup, let us prepare the handshake between our data channels and the algorithm. To do this, we need to create the sagemaker.session.s3_input objects from our data channels. These objects are then put in a simple dictionary, which the algorithm consumes.

train_data = sagemaker.inputs.TrainingInput(
    s3_train_data,
    distribution="FullyReplicated",
    content_type="text/plain",
    s3_data_type="S3Prefix",
)
data_channels = {"train": train_data}

We have our Estimator object, we have set the hyper-parameters for this object and we have our data channels linked with the algorithm. The only remaining thing to do is to train the algorithm. The following command will train the algorithm. Training the algorithm involves a few steps. Firstly, the instance that we requested while creating the Estimator classes is provisioned and is setup with the appropriate libraries. Then, the data from our channels are downloaded into the instance. Once this is done, the training job begins. The provisioning and data downloading will take some time, depending on the size of the data. Therefore it might be a few minutes before we start getting training logs for our training jobs. The data logs will also print out Spearman's Rho on some pre-selected validation datasets after the training job has executed. This metric is a proxy for the quality of the algorithm.

Once the job has finished a “Job complete” message will be printed. The trained model can be found in the S3 bucket that was setup as output_path in the estimator.

bt_model.fit(inputs=data_channels, logs=True)

Hosting / Inference

Once the training is done, we can deploy the trained model as an Amazon SageMaker real-time hosted endpoint. This will allow us to make predictions (or inference) from the model. Note that we don’t have to host on the same type of instance that we used to train. Because instance endpoints will be up and running for long, it’s advisable to choose a cheaper instance for inference.

bt_endpoint = bt_model.deploy(initial_instance_count=1, instance_type="ml.m4.xlarge")

Getting vector representations for words [including out-of-vocabulary (OOV) words]

Since, we trained with ``subwords = “True”``, we can get vector representations for any word - including misspelled words or words which were not there in the training dataset.

If we train without the subwords flag, the training will be much faster but the model won’t be able to generate vectors for OOV words. Instead, it will return a vector of zeros for such words.

Use JSON format for inference

The payload should contain a list of words with the key as “instances”. BlazingText supports content-type application/json.

from sagemaker.serializers import JSONSerializer

bt_endpoint.serializer = JSONSerializer()

words = ["awesome", "awweeesome"]
payload = {"instances": words}
response = bt_endpoint.predict(payload)
vecs = json.loads(response)
print(vecs)

As expected, we get an n-dimensional vector (where n is vector_dim as specified in hyperparameters) for each of the words.

Evaluation

We can evaluate the quality of these representations on the task of word similarity / relatedness. We do so by computing Spearman’s rank correlation coefficient (Spearman, 1904) between human judgement and the cosine similarity between the vector representations. For English, we can use the rare word dataset (RW), introduced by Luong et al. (2013).

s3.download_file("sagemaker-sample-files", "datasets/text/stanford_rare_words/rw.zip", "rw.zip")
!unzip "rw.zip"
!cut -f 1,2 rw/rw.txt | awk '{print tolower($0)}' | tr '\t' '\n' > query_words.txt

The above command downloads the RW dataset and dumps all the words for which we need vectors in query_words.txt. Let’s read this file and hit the endpoint to get the vectors in batches of 500 words to respect the 5MB limit of SageMaker hosting.

query_words = []
with open("query_words.txt") as f:
    for line in f.readlines():
        query_words.append(line.strip())

query_words = list(set(query_words))
total_words = len(query_words)
vectors = {}

import numpy as np
import math
from scipy import stats

batch_size = 500
batch_start = 0
batch_end = batch_start + batch_size
while len(vectors) != total_words:
    batch_end = min(batch_end, total_words)
    subset_words = query_words[batch_start:batch_end]
    payload = {"instances": subset_words}
    response = bt_endpoint.predict(payload)
    vecs = json.loads(response)
    for i in vecs:
        arr = np.array(i["vector"], dtype=float)
        if np.linalg.norm(arr) == 0:
            continue
        vectors[i["word"]] = arr
    batch_start += batch_size
    batch_end += batch_size

Now that we have gotten all the vectors, we can compute the Spearman’s rank correlation coefficient between human judgement and the cosine similarity between the vector representations.

mysim = []
gold = []
dropped = 0
nwords = 0


def similarity(v1, v2):
    n1 = np.linalg.norm(v1)
    n2 = np.linalg.norm(v2)
    return np.dot(v1, v2) / n1 / n2


fin = open("rw/rw.txt", "rb")
for line in fin:
    tline = line.decode("utf8").split()
    word1 = tline[0].lower()
    word2 = tline[1].lower()
    nwords += 1

    if (word1 in vectors) and (word2 in vectors):
        v1 = vectors[word1]
        v2 = vectors[word2]
        d = similarity(v1, v2)
        mysim.append(d)
        gold.append(float(tline[2]))
    else:
        dropped += 1
fin.close()

corr = stats.spearmanr(mysim, gold)
print("Correlation: %s, Dropped words: %s%%" % (corr[0] * 100, math.ceil(dropped / nwords * 100.0)))

We can expect a Correlation coefficient of ~40, which is pretty good for a small training dataset like text8. For more details, please refer to Enriching Word Vectors with Subword Information

Stop / Close the Endpoint (Optional)

Finally, we should delete the endpoint before we close the notebook.

bt_endpoint.delete_endpoint()