Building Word Embeddings with PyTorch and Lightning AI Part 1: Setting Up the Word Embedding Model

In this article, we will explore how to implement word embeddings using PyTorch and Lightning AI.

The implementation is based on the same example that we used in my Word2Vec article series.

First, let's import the required modules.

import torch
import torch.nn as nn

from torch.optim import Adam
from torch.distributions.uniform import Uniform
from torch.utils.data import TensorDataset, DataLoader

import lightning as L

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

These imports are used for the following purposes:

• torch – Create tensors and use helper functions.
• torch.nn – Create neural network layers and weights.
• Adam – Optimize the neural network using backpropagation.
• Uniform – Initialize the network weights.
• TensorDataset and DataLoader – Prepare and load the training data.
• pandas, matplotlib, and seaborn – Analyze and visualize the results.

Creating the Training Inputs

Now let's create the word embeddings for the following two sentences:

• The Incredibles is great!
• Despicable Me is great!

The first step is to create a token for every unique word in the training data and connect those tokens to the word embedding network.

Suppose we want to pass "The Incredibles" through the network.

In that case, we set its corresponding input to 1 and all the other inputs to 0.

This representation, where one input is 1 and all the others are 0, is called one-hot encoding.

In PyTorch, we can represent "The Incredibles" using a four-element vector with 1 in the first position and 0 everywhere else.

Similarly, we create one-hot vectors for the remaining words.

[
    [1., 0., 0., 0.],
    [0., 1., 0., 0.],
    [0., 0., 1., 0.],
    [0., 0., 0., 1.]
]

Finally, we convert these vectors into a PyTorch tensor.

inputs = torch.tensor([
    [1., 0., 0., 0.],
    [0., 1., 0., 0.],
    [0., 0., 1., 0.],
    [0., 0., 0., 1.]
])

Now that we have created the input data, the next step is to define the labels. We will explore that in the next article.

AI agents write code fast. They also silently remove logic, change behavior, and introduce bugs -- without telling you. You often find out in production.

git-lrc fixes this. It hooks into git commit and reviews every diff before it lands. 60-second setup. Completely free.

Any feedback or contributors are welcome! It's online, source-available, and ready for anyone to use.

Give it a ⭐ star on Github