Get Neural Network to predict a tag/class on a certain word using the surrounding words as context [PyTorch]?
You may want to take a look at this article, but I'll summarize. You can use BERT (or some other tool) to make embeddings of every word in every sentence. Then for each word, make a contextualized embedding vector using the rest of the sentence. bert-embedding does all of this itself. Then keep the embedding vector for the important words.
For each important word, you would then have two pieces of information: the embedding vector and the correct label (which could easily be made into an integer from $0$ to $4$). Depending on the size of the embedding vectors, you could use PCA to reduce the size, although this may not be needed. Using this data, you can then train a neural network or use a k-nearest neighbors classifier.
There is more information in the article, which I suggest that you read. They do a better job explaining than me and also have some actual code you may want to look at.