Word Embedding in PyTorch + Lightning

Key Takeaways

[Music] wood edding with po torch and lightning hooray stack Quest hello I'm Josh starmer and welcome to stack Quest today we're going to talk about word embedding in pie torch plus lightning don't stress out about the cloud use lightning bam this stack Quest has also brought to by the letters a b and c a always b b c curious always B curious note this stack Quest assumes you are already familiar with word embedding if not check out the quest also note you can download all of the code in this stack quest for free the details are in the pinned comment below in the stack Quest on word embedding we created a simple neural network that converted the word or input tokens Troll 2 is great and Jim cata into numbers which we call word embeddings we also showed that these word embeddings allow words that are used in similar contexts like Troll 2 and Jim Kata to appear close to each other when we use the embedding values to plot each word on a graph bam now in this stack Quest we'll learn how to build and train this simple word embedding network with P torch plus lightning first we'll do it from scratch using just tensors and some basic math and then we'll simplify our code using the pi torch linear function lastly we'll learn how to use the pytorch embedding function to load and use pre-trained word embeddings the first thing we do is import torch to create the tensors we will use to store the raw data and to provide a few helper functions then we import torch.nn to create the weights we will use in the network and to bring in some other helper functions then we import atom to fit the neural network to the data with back propagation then we import uniform to initialize the weights in the network and to give us the tools to create a large scale embedding network with lots of training data we'll import tensor data set and data loader from torch. ts. data now we import lightning as L to make it way easier to write our code and for automatic code optimization and scaling in the cloud lastly we import pandas matte plot lib and Seaborn so that we can draw some pretty graphs now we will create word embeddings for these two simple sentences Troll 2 is great and Jim Kata is great so we will need to create an input for for each unique token in the training data and these inputs will eventually connect to this word embedding Network now if we wanted to run Troll 2 through the network we would put a one in the input for Troll 2 and a zero in all of the other inputs oh no it's the dreaded terminology alert when we specify the inputs like this where one input gets a one and everything else gets a zero it's called one hot encoding so if we want to run the word is through the network we specify the one hot en coding with a one for is and a zero for everything else likewise this is the one hot en coding for great and this is the one hot en coding for Jim cata now going back to the one hot en coding for trol 2 in pi torch we can create a Four Element list with one in the the first position for Troll 2 and zeros in all of the other positions to recreate the one hot encoding for Troll 2 likewise we can create similar one hot encoding lists for is great and Jim Kata and once we have a list for each possible input we make a list of these lists and because we're using pi torch we convert the input lists into tensors with torch. tensor and save the tensors in a variable called inputs now the goal is to create a simple Network that can predict the token that follows a specific input for example if the input token is Troll 2 then we want to predict the word is and that means we want the output for is to be one and the outputs for all of the other tokens to be zero in pi torch that means we want to predict the one hot encoding for is likewise when the input is is then we want to predict great and thus we want to predict the one hot encoding for great great is at the end of each phrase and nothing comes after it so in theory it shouldn't predict anything but if we had a larger training data set it probably would predict something so in this example will just pretend that it predicts Jim Kata lastly Jim Kata just like Troll 2 predicts is now let's combine the output lists in a list and convert the output lists into tensors and save them in a variable called labels because in machine learning that is what we call the known or ideal output values and now we are done encoding the training data now that we have the inputs and label we can combine them into a tensor data set that we'll call data set and then use data set to create a data loader called Data loader data loaders are super useful when we have a lot of data because one they make it easy to access the data in batches two they make it easy to shuffle the data each Epoch and three they make it easy to use a relatively small fraction of the data if we want to do a quick dirty training for debugging but Josh we don't have a lot of data why are we using a data loader you're right Squatch we don't have a lot of data but in a more realistic setting we would so we might as well do it now okay anyway now that we have the data taken care of let's write the code for this simple word embedding Network when we create a neural network in pi torch we always start by defining a new new class now because we're coding our first word embedding network from scratch we'll call this word embedding from scratch and in order to make coding super easy We'll Inherit from lightning module then just like we always do we create an initialization method for the new class this method will create and initialize all of the weight tensors that we need to implement the embedding Network and it will also create the LW function that we'll use during training then we'll create a method called forward that makes a forward pass through the embedding Network then we'll create a method to configure the atom Optimizer and lastly we'll create a method called training step to calculate the loss which in this case will be the cross entropy loss the cross entropy loss function will quantify the difference between what we want the out output to be and what we actually get for the output note there's a lot more to be said about the cross entropy loss function so if you're curious check out the quest now let's start by coding the anit method the first thing we do is call the initialization method for the parent class lightning module this is simply required whenever we inherit from a class in Python in this case this will allow us to take advantage of all the features that lightning offers now we need to create and initialize the weights for the network and we're going to do this by using a uniform distribution to randomly select an initialization value for each weight specifically we're going to use this uniform distribution that goes from 0.5 to 0.5 to Generate random numbers for the weights the shape of this distribution shows the that all values between 0.5 and 0.5 have the same likelihood of getting randomly selected hey Josh why are we using values between 0.5 and 0.5 to initialize the weights good question Squatch we're using this specific range of values in order to match up with what we will do in the second part of this tutorial when we use the pi torch linear function to do the math and the L linear function selects a range of values based on the number of inputs okay anyway in order to use this uniform distribution to Generate random numbers for us we create a variable called Min value and set it to 0.5 the minimum value we want to randomly select and we create a variable called max value and set it to 0.5 the maximum value we want to randomly select now we create a parameter for the first weight associated with the first input and use uniform. sample which we imported earlier to initialize it with a random number when I did this the first weight associated with Troll 2 was randomly set to 0.38 then the second weight associated with Troll 2 was randomly set to 0.42 now we just create an initialize parameters for all of the other weights associated with each input bam now we've created and initialized all of the weights associated with the inputs likewise we create and initialize parameters for all of the weights associated with each output now we have all of the weights initialized for our word embedding Network bam now the last thing we need to do in our anit method is give our class access access to the Cross entropy loss function and we do this by calling nn. cross entropy loss and saving it in a variable called loss now that we're done coding the anit method we can use those weights to code the forward method to make a forward pass through the embedding Network for the forward method the input is a list that contains the one hot encoding for one of the input tokens for example the input might be the one hot encoding for Troll 2 however when it is passed to the forward method it comes wrapped up in an extra set of brackets so the first thing we do is remove those brackets by setting input to be the first element now we multiply each input value by its corresponding weight that goes to the activation function on top and add the products together and we save that s in a variable called inputs to top hidden then we multiply the inputs by the weights to the activation function on the bottom and add the products together and save the sum in a variable called inputs to bottom hidden so now we have the code for the first part of the word embedding Network and since the activation functions are identity functions which means the input is the same as the output we can multiply inputs to top hidden and inputs to bottom hidden by the next set of Weights directly for example we can multiply inputs to top hidden by the first weight going to the top output and then multiply inputs to bottom hidden by the second weight going to the top output and then add those two products together and save the result in a variable called output one then we do the same same thing for the other outputs bam now we've done all the math up to the softmax function and that actually means we are done making a forward pass through the embedding Network because the loss function that we're using for back propagation nn. crossentropy loss does the soft Max for us so the last thing we need to do is package up the output values using torch. stack and save everything in a variable called output pre- soft Max note if instead of using torch. stack we just returned a list of the output values by wrapping them up in square brackets then the gradients would get stripped off and we would not be able to do back propagation so by using torch. stack we can return a list that preserves the gradients anyway the last thing we do in the forward method is return output pre- soft Max now that we have the forward method we are ready to configure the optimizer and configuring the optimizer in this case atom is so easy we can just replace the pseudo code with the real code we pass atom the parameters we want to optimize and we set the learning rate LR to 0.1 hey Josh why did you set the learning rate to 0.1 because our example is pretty simple and I wanted to train relatively quickly I tested out a relatively large learning rate 0.1 and it worked okay now let's talk about the training step method which we'll use to calculate the loss the training step method takes a batch of training data and the index for that batch and the first thing we do is split the batch of training data into the input and the labels which are the ideal output values then we run the input through the network up to the softmax function by passing it to the forward method for example if we run the one hot encoding for Troll 2 through the untrained Network these are the values that the forward method will return we then run those values along with the ideal values through the loss function nn. cross loss then runs the output values through a soft Max function and quantifies the difference between the soft Max output in the ideal values and we save that difference in a variable called loss and then return the loss now at long last we've made it through all the code needed to create word embeddings from scratch we create and initialize the weight tensors and create the loss function in the init method we make a forward pass through the embedding network with the forward method configure the atom Optimizer with configure uncore optimizers and last but not least calculate the loss with training step bam now let's use the new class we just wrote to create a new word embedding Network that we'll call model from scratch and let's print out the randomly selected weight values that it starts out with here we just have a for Loop that iterates over all the named parameters in the network and for each parameter it prints out its name and value and here's the output now to be honest this list of numbers is kind of hard to read so let's organize it into an easyto read data frame so the first thing we do is put the weight values into a dictionary the first first part label W1 contains the weight values for each input that goes to the activation function on top note we're using the item method to get the weights because it Returns the tensor values as python numbers the second part label W2 contains the weight values for each input that goes to the activation function on the bottom then we just label the tokens and inputs and save the dictionary in the very aable called data and then transform data into a pandis data frame called DF and print out DF and this is what the data frame looks like now we can easily see that the weights for Troll 2 and Jim Kata are relatively different even though they both represent movie titles that are used in the same context this table is pretty helpful for making the embedding values easy to look at but a graph would make them even easier to look at this graph has the weight values to the top activation function W1 on the x axis and the weight values to the bottom activation function W2 on the Y AIS with a graph it's super easy to see that the embedding values for troll two are very different from the values for Jim Kata to create the graph the the first thing we do is call the caborn function scatter plot and we pass scatter plot the data frame DF that we just created and we tell scatter plot that we want to use the weights that go to the top activation function W1 on the X AIS and the weights that go to the bottom activation function W2 on the Y AIS now if all we did was call scatter plot then we'd end up with this scatter plot and while this scatter plot is super cool it would be much cooler if each dot were labeled with the word or token that it represented like this so in order to add the tokens as labels to each point we call the Matt plot lib text function and we pass in the X and Y AIS coordinates for the point in the first row in the data frame and the value for the token then then we specify how we want the text aligned the font size the font color and lastly the font weight then we do the same thing for each row in the data frame and then we call PLT doow and we get this super cool looking scatter plot like we mentioned earlier we can now see that the embedding values for Troll 2 and Jim cata are pretty different and that means we need to train our embedding Network we start training by creating a lightning trainer called trainer and tell it to train for at most 100 Epoch which means we will do back propagation for every weight using the training data at most 100 times now we call the trainers fit method and pass it the embedding Network called model from scratch and the training data called Data loader in theory it should only take a few seconds to train our simple embedding Network and when it's done we recreate the data frame that has the weights or embedding values for each token and we can either stare at the embedding values in the data frame or we can draw a scatter plot of the tokens just like before note because the labels for Troll 2 and Jim cata are overlapping and hard to read I then added little offsets to where Troll 2 and Jim cata were printed and now we can see that after training the embedding Network the embedding values for Troll 2 and Jim cata are very similar which is great since they are used in similar contexts bam now that we have trained our embedding Network we can see what it predicts when we use Troll 2 as the input remember from when we created the training data that we want Troll 2 to predict is so the first thing we need to do is create a soft Max function because we didn't have to explicitly use it in our model note we set dim equal to zero so that we can apply it to rows of output values if we set dim equal to one then we would apply it to Columns of values now we pass Troll 2 as a one hot encoded tensor into model from scratch and we run the output values through the the soft Max and then round the output of the soft Max to two decimal places and finally print out the result and we get the one hot encoded tensor for is which is correct bam likewise we can verify that all of the other inputs to our embedding Network create the correct outputs okay now that we know how to create and train a simple word embedding Network from scratch let's make our lives a little easier by using the P torch linear function to create the same network so let's create a new class called word embedding with linear and again in order to make training super easy will inherit from lightning module then create the anit method that we use to create and initialize the weights and like always we'll call the anit method from the parent class now comes the interesting part instead of calling nn. parameter to create and initialize each weight in the network we only have to make two calls to nn. linear the first call creates the weights between the inputs and the hidden layer in features equals 4 means we are connecting four inputs to two nodes specified with features equals 2 in the hidden layer in other words this call to nn. linear will make four weights for each of the two nodes in the hidden layer and since we don't need any bias terms we set bias equal to false the second call to nn. linear creates the weights between the hidden layer and the outputs it creates two weights within features equal to two for each of the four outputs without features equal to 4 and again since we don't need any bias terms we set bias equal to false now the last thing we need to do in our anit method is give our class access to the Cross entropy loss function now we need to code the forward method that makes a forward pass through the network the cool thing is that all we have to do to calculate the sums before the activation functions is pass the input to the linear object input to Hidden that we created in the anit method and save the sums in a variable called hidden the linear object input to Hidden does all of the multiplication and addition for us bam note now that we are using linear to do the math we no longer have to strip off the extra brackets from input like we did before anyway because the input to these activation functions is the the same as the output we can just ignore them and pass hidden to the second linear object we created hidden to Output hidden to Output calculates the output sums from the activation functions and we save those output values in output values now remember that we don't need to calculate the soft Max because the loss function nn. crossentropy loss does it for us so all we have to do is return the output values bam the next thing we do is create the configure optimizers method and just like before we use the atom Optimizer and pass it the parameters we want to optimize and set the learning rate to 0.1 the training step method is also similar to what we did earlier and that means that this embedding network is contained in this class definition by using inn. linear we significantly reduced the amount of code we need compared to when we did everything from scratch we can see the difference when we shrink the original code for the word embedding from scratch class down so that we can fit it on a single screen and compare it to the size of the word embedding with linear class on the same hardto read scale bam now we can create a new model model linear with our new class word embedding with linear and just like before put the pre-trained word embedding values into a data frame called DF note because we used nn. linear to create the weights we access them with weight and we call detach to remove the gradient from the tensors and we use zero and one to index the weights that go to the top and bottom activation functions and lastly convert the tensor to a num Pi array with num Pi now when we print out our data frame DF we get this nicely formatted table and we can draw a scatter plot of the tokens just like before and that gives us a scatter plot that looks like this and the graph suggests that the embedding values or weights are not yet optimal because Troll 2 and Jim are so far from each other so just like before we can train the model for 100 Epoch and after training we end up with these weights and the weights going from the inputs to the hidden layer are the new embedding values and when we redraw the scatter plot with the new embedding values we see that the values for Troll 2 and Jim cata are similar double bam now that we we know how to create word embedding networks from scratch and within in. linear and we can create word embeddings that put words and tokens used in similar context near each other let's learn how we can load and use pre-trained word embedding values with nn. edding Note in this example we're just going to load and access the embedding values we created but in practice we might want to load and use the word Tove embedding values which have 100 values per token and millions of tokens or the embedding values created by a Transformer like chat GPT however before we get started let's just print out the embedding values from the last Model we trained model linear we access the embedding values in model linear just like we did when we created the data frame except now we don't have to worry about the gradient an and we get two lists of embedding values specifically the weights are arranged in two rows the first row corresponds to the weights that go to the top activation function and the second row corresponds to the weights that go to the bottom activation function now the problem with this is that nn. embedding expects the weights to be in columns just like in the data frames we created the good news is that converting rows into columns is super easy so with that said let's load and use these pre-trained word embedding values with nn. edding we start by creating an nn. embedding object and we pass the pre-trained weights in with from pre-trained and we use T to transpose the rows of Weights into columns lastly we save the the new embedding object in word embeddings we can then verify that we did things correctly by printing out the weights and we see that the weights are now arranged in two columns now we can print out the embedding values for the first input Troll 2 by passing in a tensor with the first index value zero and the embedding values match what we expect so we know we did things correctly accessing the embedding values by index is fine but we can also make our lives Easier by creating a dictionary that maps The Tokens to their indices and now we can more easily access the embeddings with the token itself rather than the index and that's all there is to loading and accessing pre-trained weights into an nn. embedding object we can now use our embedding object word embeddings and connect it to a larger neural network like a Transformer also before we go I just want to remind you that you don't have to type this code yourself instead you can download it and I wrote tons of comments that explain every little detail just like this stack Quest the link is an app pinned comment below triple bam now it's time for some Shameless self-promotion if you want to review statistics and machine learning offline check out the stack Quest PDF study guides and my book the stat Quest Illustrated guide to machine learning at stat quest.org there's something for everyone hooray we've made it to the end of another exciting stack Quest if you like this stack Quest and want to see more please subscribe and if you want to support stack Quest consider contributing to my patreon campaign becoming a channel member buying one or two of my original songs or a t-shirt or a hoodie or just donate the links are in the description below all right until next time Quest on