PySparkでWord2Vecを実行する

はじめに

前回に続いて、Spark（入りコンテナ）を使ってみました。

最近、大規模言語モデルが世間を騒がせているので、その関連で個人的に馴染み深いWord2VecモデルのSparkでの扱い方をご紹介。
本家の掲載Codeを参考に、Wikipedia日本語全文を利用してWord2Vecモデルを作り、類似単語の検索に使いました。

本稿で紹介すること

• 環境の各種情報
• Spark入りDockerコンテナの起動
• Word2Vecの実行

紹介するPythonコードはリファクタリング未済のため、PEP8には則っていません！

本稿で紹介しないこと

• WSL2のインストール
• DockerCEのインストール
• JupyterHubのインストール＆環境設定
• Word2Vecの全般

環境の各種情報

筆者は、Windows 11のホスト上でWSL2／Ubuntuを起動しています。

• Ubuntu 20.04.6 LTS
• Docker CE 23.0.1

更に筆者は、WSL2／Ubuntuのホスト上でSpark入りDockerコンテナを起動しています。

• Ubuntu 22.04.1 LTS
• Python 3.8.13
• pip 22.2.2
• pyspark 3.3.0
• OpenJDK 17.0.4
• Spark 3.3.0

コンテナイメージはコチラ。

以下コマンドを実行してコンテナイメージを取得可能です。

$ docker pull jupyter/pyspark-notebook:python-3.8

Spark入りDockerコンテナの起動

JupyterHubからSpark入りDockerコンテナを起動しました。
が、以下コマンドを実行してDockerコンテナを起動することも当然可能です。

$ docker run -d -it --name spark -p 8888:8888 jupyter/pyspark-notebook:python-3.8

Word2Vecの実行

Spark公式で公開されている例¹²を利用して実行しました。
ざっくり言うと、分かち書きからDataFrameに変換する部分を追加しました。

以下、Spark公式で公開されている例です。
ポイントは、pyspark.ml.feature.Word2Vecを利用することです。

from pyspark.ml.feature import Word2Vec

# Input data: Each row is a bag of words from a sentence or document.
documentDF = spark.createDataFrame([
    ("Hi I heard about Spark".split(" "), ),
    ("I wish Java could use case classes".split(" "), ),
    ("Logistic regression models are neat".split(" "), )
], ["text"])

# Learn a mapping from words to Vectors.
word2Vec = Word2Vec(vectorSize=3, minCount=0, inputCol="text", outputCol="result")
model = word2Vec.fit(documentDF)

result = model.transform(documentDF)
for row in result.collect():
    text, vector = row
    print("Text: [%s] => \nVector: %s\n" % (", ".join(text), str(vector)))

Word2Vecモデルを作る

過去記事を参考に、Wikipedia日本語全文³⁴からInputデータ、分かち書き（Code中のjawiki_wakati.txt）を予め作っています。

# $example on$
from pyspark.ml.feature import Word2Vec, Word2VecModel
# $example off$
from pyspark.sql import SparkSession

sentences = [(s.split(" "), ) for s in open("jawiki_wakati.txt", "r").readlines()[:10000]]

spark = SparkSession.builder \
    .appName("word2vec-spark") \
    .getOrCreate()

# Input data: Each row is a bag of words from a sentence or document.
documentDF = spark.createDataFrame(sentences, ["text"])

# Learn a mapping from words to Vectors.
#word2Vec = Word2Vec(vectorSize=100, minCount=0, inputCol="text", outputCol="result")
word2Vec = Word2Vec(vectorSize=100, minCount=5, windowSize=5, numPartitions=1, maxIter=1, seed=777, inputCol="text", outputCol="result")
model = word2Vec.fit(documentDF)

#model.save(spark, 'w2v_model')
model.save('w2v_model')

spark.stop()

Word2Vecモデルを使う

近いベクトルの単語≒類似単語を検索しています。

# $example on$
from pyspark.ml.feature import Word2Vec, Word2VecModel
# $example off$
from pyspark.sql import SparkSession
from pyspark.sql.functions import format_number as fmt

spark = SparkSession.builder \
    .appName("word2vec-spark") \
    .getOrCreate()

#model = Word2VecModel.load(sc, 'w2v_model')
model = Word2VecModel.load('w2v_model')

# 近いベクトルの単語を取得
model.findSynonymsArray('野球', 5)

spark.stop()

いろいろな実行結果の確認

筆者の環境（マシンSpec）と処理時間を鑑みて、Inputデータを制限しているので、実行結果はイマイチな精度な印象です。
が、Inputデータを増やすことで実行結果の精度、納得感のある類似単語を取得可能になることも確認済です。

# 近いベクトルの単語を取得
model.findSynonymsArray("東京", 5)

# 実行結果
[('バス', 0.9257544875144958),
 ('上野', 0.9251026511192322),
 ('都立', 0.9039059281349182),
 ('京成', 0.9008254408836365),
 ('清', 0.8926427364349365)]

# 近いベクトルの単語を取得
model.findSynonymsArray('日本', 5)

# 実行結果
[('政府', 0.7264666557312012),
 ('ウィーン', 0.6764424443244934),
 ('クラブ', 0.6741369366645813),
 ('トルコ', 0.6721512079238892),
 ('都市', 0.6598173379898071)]

ちなみに、以下のようなCodeでも同様の結果を得られます。

model.findSynonyms('日本', 5).select("word", fmt("similarity", 5).alias("similarity")).show()

# 実行結果
+--------+----------+
|    word|similarity|
+--------+----------+
|    政府|   0.72647|
|ウィーン|   0.67644|
|  クラブ|   0.67414|
|  トルコ|   0.67215|
|    都市|   0.65982|
+--------+----------+

Notebook（Pythonコード）

GitHubで公開しています。

まとめ

前回に同じく、完全に手探りでしたが、Sparkを使ってK-meansクラスタリングを実行する方法を紹介しました。
また、小規模データならばSparkである必要性やありがたみを感じれませんが、大規模データならばその限りではないと思いました。

1. Feature Extractors | https://spark.apache.org/docs/3.3.0/ml-features#word2vec ↩

2. word2vec_example.py | https://github.com/apache/spark/blob/master/examples/src/main/python/ml/word2vec_example.py ↩

3. Wikipedia:データベースダウンロード | https://ja.wikipedia.org/wiki/Wikipedia:%E3%83%87%E3%83%BC%E3%82%BF%E3%83%99%E3%83%BC%E3%82%B9%E3%83%80%E3%82%A6%E3%83%B3%E3%83%AD%E3%83%BC%E3%83%89 ↩

4. Index of /jawiki/latest/ | https://dumps.wikimedia.org/jawiki/latest/ ↩