Word embedding using Scala and Spark (part 1)

In this post, I will show a very simple how to calculate the Word Embeddings using the SparkMLlib in Scala. In this post, I will reuse the well-known technique Word2Vec.
Word2Vec is a group of techniques that allow you to takes a large corpus of text and produces for each word a vector space (typically of hundred of dimensions).  This prototype is formed by 3 modules:

  • Document repository
  • Input module
  • Word2Vec module

The repository could be a database, a web services, but in this case, it is a simple text file, called data.dat, where there is a document for each row. For this demo, the file contains just 4 documents:

  • The sun is a star
  • The earth is a planet
  • The moon is a satellite’s earth
  • The sun is yellow

Before to describe the reading module, let’s see the libraries used for this job.

import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.ml.feature.{Word2Vec, Word2VecModel}
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
val spark2 = spark
import spark2.implicits._

Over the typical Spark libs, we will use some org.apache.spark.ml libs like Word2Vec/Word2VecModel for calculating embeddings. Now, we can describe the reading module. It is very simple, in fact, it is the line for the creation of a DataFrame from the data.dat file:

val documents = sc.textFile("../files/data.dat").
            map(line => line.toLowerCase).
            map(_.split(" ")).

More in details, for each row, we will perform the toLowerCase transformation and then the splitting of each document into words. The result is a DataFrame with a column called text, where each row is the array contains the split document:

  • [the, sun, is, a, star]
  • [the, earth, is, a, planet]
  • [the, moon, is, a, satellite’s, earth]
  • [the, sun, is, yellow]

After, we have to create the Word2Vec model. The code is the following:

val word2vec = new org.apache.spark.ml.feature.Word2Vec().

First, we use the Word2Vec class for creating an object that takes the documents Dataframe and produces the embeddings in the column “feature”. For each word, the module will create a vector of 100 embeddings.
Moreover, we can set the minimum frequency a token for being included in the word2vec model’s vocabulary (setMinCount) and the max number of iterations (setMaxIter).
For seeing the vector for the word sun, we can use the following:

modelW2V.getVectors.filter($”word” === “sun”).show()

The vector with 100 elements for the word “sun” is:


Now, we can train the model with the opened documents:

val modelW2V = word2vec.fit(documents)

Due to the limited number of documents, the model is ready in a few seconds. Next step is the creation of the embeddings for every document. This is the first problem. How can I combine the word’s vector for a document? Without losing generalization, we can think that the vector for a document is a combination (max, min, average) of embeddings for each word. For this, SparkML helps us with the methods transform.
In fact, the transform method, that belongs to the Word2VecModel class, use the average function for creating the vector for a document.

// Creation of embeddings for documents
val parsedData = modelW2V.transform(documents)

With less than 20 lines of codes, we have the vector for each word and document.

For now, that’s all…


Sharing is caring!

Leave a Reply