这是pyspark 中的一个示例,我想它可以直接移植到 Scala - 关键是model.transform 的使用。
首先,我们按照示例训练模型:
from pyspark import SparkContext
from pyspark.mllib.feature import Word2Vec
sc = SparkContext()
inp = sc.textFile("text8_lines").map(lambda row: row.split(" "))
k = 220 # vector dimensionality
word2vec = Word2Vec().setVectorSize(k)
model = word2vec.fit(inp)
k 是词向量的维数 - 越高越好(默认值为 100),但您需要内存,我的机器可以使用的最高数字是 220。(编辑: 相关出版物中的典型值在 300 到 1000 之间)
在我们训练好模型之后,我们可以定义一个简单的函数如下:
def getAnalogy(s, model):
qry = model.transform(s[0]) - model.transform(s[1]) - model.transform(s[2])
res = model.findSynonyms((-1)*qry,5) # return 5 "synonyms"
res = [x[0] for x in res]
for k in range(0,3):
if s[k] in res:
res.remove(s[k])
return res[0]
现在,这里有一些国家和首都的例子:
s = ('france', 'paris', 'portugal')
getAnalogy(s, model)
# u'lisbon'
s = ('china', 'beijing', 'russia')
getAnalogy(s, model)
# u'moscow'
s = ('spain', 'madrid', 'greece')
getAnalogy(s, model)
# u'athens'
s = ('germany', 'berlin', 'portugal')
getAnalogy(s, model)
# u'lisbon'
s = ('japan', 'tokyo', 'sweden')
getAnalogy(s, model)
# u'stockholm'
s = ('finland', 'helsinki', 'iran')
getAnalogy(s, model)
# u'tehran'
s = ('egypt', 'cairo', 'finland')
getAnalogy(s, model)
# u'helsinki'
结果并不总是正确的 - 我将留给您进行实验,但随着更多的训练数据和增加的向量维度k,它们会变得更好。
函数中的for 循环删除属于输入查询本身的条目,因为我注意到通常正确答案是返回列表中的第二个,而第一个通常是输入项之一。