用 Apache Spark 來訓練 Word2vec 詞向量 Skip-gram Word Embeddings

先前曾提過 word2vec 可以把語料中的詞轉換成詞向量。雖然原本提供的 word2vec 工具速度已經很快，但是如果要訓練更大規模的語料還是需要不少時間。例如之前在處理 ClueWeb09 時，以實驗室的機器來說，就算只處理中文部份，也要一個月以上才能跑完。此時除了購買更強大的機器外，如果已經有不少機器，或許可以利用平行運算的方式來加速。

最近剛好接觸到了 Apache Spark，他是一個開源的運算平台，可以運用多台電腦進行平行運算。且因為把很多資料直接放在記憶體中處理，又比 Apache Hadoop 單純的 MapReduce 更快一些。更重要的是，他的機器學習函式庫 MLlib 已經實作了 word2vec 當中的 skip-gram 模型，正好可以直接拿來訓練詞向量。

安裝

關於如何將 Spark 安裝在一個 cluster 上，可以參考官方文件。這裡我們只簡單的安裝單機版的 Spark 方便快速的實驗。我們將會使用 Ubuntu 14.04 作為實驗平台。

安裝 Java 8

首先安裝 Java 8，如果你已經有裝了則可以跳過。

sudo add-apt-repository ppa:webupd8team/java
sudo apt-get update
sudo apt-get install oracle-java8-installer
sudo apt-get install oracle-java8-set-default

安裝 Apache Spark

接下來，到下載頁面下載 Spark，我是選擇 Spark 1.4.1 Pre-built for Hadoop 2.6 and later 的 binary。不過如果有新版應該變化也不大。直接解壓縮就可以用囉：

tar -xzf spark-1.4.1-bin-hadoop2.6.tgz

安裝 sbt

因為我打算用 Scala 所以我們還得安裝 sbt 這個 build tool。於是依照 sbt 的下載頁面的說明安裝 sbt：

echo "deb https://dl.bintray.com/sbt/debian /" | sudo tee -a /etc/apt/sources.list.d/sbt.list
sudo apt-key adv --keyserver hkp://keyserver.ubuntu.com:80 --recv 642AC823
sudo apt-get update
sudo apt-get install sbt

下載資料集

為了簡單起見，使用 100MB 的 Wikipedia dump：

wget http://mattmahoney.net/dc/text8.zip -O text8.gz
gzip -d text8.gz -f

由於是單機，所以可以放在任意資料夾，如果是在 cluster 的話，需要放在每台機器都能存取的同一位置。

編寫 Spark 程式

首先建立專案資料夾：

mkdir -p sparkw2v/src/main/scala/

然後編輯 sparkw2v/sparkw2v.sbt 檔案，程式版本參考官方文件：

name := "Spark Word2Vec"

version := "1.0"

scalaVersion := "2.10.4"

libraryDependencies ++= Seq(
    "org.apache.spark" %% "spark-core" % "1.4.1" % "provided",
    "org.apache.spark" %% "spark-mllib" % "1.4.1"
)

最後則是程式碼 sparkw2v/src/main/scala/SparkW2V.scala 本身，注意要設定輸出入 {input directory} 和 {output directory}的路徑。同樣的如果是在 cluster 的話，需要放在每台機器都能存取的同一位置。 Word2Vec 參數設定則可參考 API 文件。

import org.apache.spark._
import org.apache.spark.rdd._
import org.apache.spark.SparkContext
import org.apache.spark.SparkContext._
import org.apache.spark.SparkConf

import org.apache.spark.mllib.feature.Word2Vec

object SparkW2V {
  def main(args: Array[String]) {
    val text8 = "{input directory}/text8"
    val output = "{output directory}/model"
    val conf = new SparkConf().setAppName("Spark Word2Vec")
    val sc = new SparkContext(conf)

    val input = sc.textFile(text8).map(line => line.split(" ").toSeq)
    val word2vec = new Word2Vec()

    val model = word2vec.fit(input)
    model.save(sc, output)
  }
}

Build

使用 sbt 進行 package：

cd sparkw2v
sbt package

根據版本不同，產生的檔案名稱也不同，我的輸出是 sparkw2v/target/scala-2.10/spark-word2vec_2.10-1.0.jar。

執行

接下來就可以回到根目錄，利用 spark-submit 執行程式。筆者不確定記憶體的需求為何，所以開的大了些：

cd ..
spark-1.4.1-bin-hadoop2.6/bin/spark-submit --class SparkW2V --master local[*] --executor-memory 20G --driver-memory 10G sparkw2v/target/scala-2.10/spark-word2vec_2.10-1.0.jar

輸出的向量會放在 {output directory}/model/data/ 底下，而且是存成 Parquet 的格式，不太方便。所以我們使用 spark-shell 快速的將檔案轉成文字檔：

# 執行 Spark Shell
./spark-1.4.1-bin-hadoop2.6/bin/spark-shell

# 從這裡開始是 Spark Shell
#
# Welcome to
#       ____              __
#      / __/__  ___ _____/ /__
#     _\ \/ _ \/ _ `/ __/  '_/
#    /___/ .__/\_,_/_/ /_/\_\   version 1.4.1
#       /_/
#

# 讀取檔案
scala> val d = sqlContext.read.parquet("{output directory}/model/data")
d: org.apache.spark.sql.DataFrame = [word: string, vector: array<float>]

# 檢查格式
scala> d.first
res2: org.apache.spark.sql.Row = [latifolia,ArrayBuffer(-0.08103186, 0.14688604, -0.060668133, -0.25648367, -0.06855837, -0...

# 輸出到 output directory/vectors/
scala> d.map{r => r.getString(0) + " " + r.getSeq(1).mkString(" ")}.saveAsTextFile("{output directory}/vectors")

如此便完成了！

原始碼

相關程式碼放在 shaform/experiments/spark_word2vec。