pyspark – 基于word2vec+LSH实现相似内容查找

本文基于kaggle豆瓣影评数据集,演示如何利用pyspark的word2vec和LSH库实现相似影评的计算,同样的方式可以用于相似内容匹配,例如:在海量文章中检测存在抄袭的文章等类似需求。

代码地址:https://github.com/owenliang/douban-comments-similarity,豆瓣影评数据集加载如下:

每行是1条影评,其关联了电影标题、影评作者、评论内容等信息,我们基于该数据集利用机器学习算法挖掘出Comment内容相似的影评。

第1步:jieba分词

下面代码对所有影评Comment进行jieba分词:

分词采用udf自定义函数对Comment列执行分布式计算,利用jieba库中文分词,并擦除掉标点符号,最终生成分词列表到Words列。(这块分词处理比较粗糙,还可以进一步优化)

第2步:word2vec训练词embedding向量

每条评论的分词序列作为样本交给word2vec模型训练,可以为每个词学得K维的embedding向量表示,两个词embedding向量之间的距离可以体现出词之间的相似度。

将每条内容的所有分词的embedding向量求平均,可以得到内容的embedding向量,相似内容之间的embedding向量距离也是相近的。

下面是word2vec原理,我们理解模型的输入输出是什么、以及embedding向量来自于模型的哪个部分即可:

  • 通过对分词序列滑动窗口,可以不断生成训练样本,中心词为输入,周边词为输出,均为onehot方式。

利用双层神经网络,训练后权重矩阵W的每一行就视作对应onehot位置单词的embedding向量了。

下面的代码基于每行评论的词序列生成大量样本,完成了word2vec训练并得到了词embedding向量:

这里我对训练好的word2vec神经网络做了保存与加载,避免下次启动重复计算。

这里指定:

  • vectorSize:词embedding向量长度。
  • numPartitions:计算并行度。
  • maxIter:样本迭代训练次数。
  • seed:随机数种子。
  • inputCol:指定Words列为词序列字段。
  • outputCol:指定内容embedding向量(也就是分词embedding的平均)的输出字段。

这里getVectors可以查看模型学到的所有word的embedding向量。

第3步:word2vec生成内容embedding向量

调用上述word2vec模型,可以直接为每条评论生成内容embedding向量,其原理为计算其所有分词embedding向量的平均值。

评论embedding向量输出在Embedding列,此时任意两条评论之间的距离可以直接基于embedding列的向量求距离(例如余弦距离、欧式距离等)即可。

为什么距离有意义呢,很简单:

第4步:训练LSH

虽然有了评论embedding向量,但是要为每条评论计算与其embedding距离最近的其他评论,需要进行笛卡尔乘积规模的全交叉运算,这个代价在大规模数据集下不现实。

LSH局部敏感哈希算法可以解决这个问题,它通过对所有embedding向量的观察,可以学得向量v:

对于embedding向量x,通过与 v内积可以映射为整数h(v),即: h(v)=v⋅x

此时对整数h(v)进行哈希分桶,LSH可以基本保证相同桶内embedding向量距离相近,而不同桶之间的embedding向量相远,因此对于任意x向量通过通过快速分桶并仅对桶内的embedding向量求距离,那么可以极大降低计算规模。

我们执行如下代码训练LSH,输入是所有的内容embedding向量,训练后的LSH模型具备了向量分桶能力:

  • numHashTables:支持多个哈希函数,这样可以对同一个embedding向量分到多个桶,每个桶内都计算与其他向量的距离,这样可以提升搜索的准确率(因为相近的embedding有概率被划入2个桶,如果只使用1个哈希则会错过其搜索机会)
  • bucketLength:能够决定分桶的数量,桶越多单个桶内的向量越少计算越快,但是误分桶的可能性越高,具体参考:https://george-jen.gitbook.io/data-science-and-apache-spark/locality-sensitive-hashing
  • outputCol:embedding向量的分桶结果,有几个哈希函数就会落入几个桶内,搜索会在多个桶内进行。

第5步:LSH求embedding的最近距离向量

因为LSH计算好了每个内容的分桶,相当于内容被投入到对应的桶内了,因此桶内的embedding之间可以俩俩计算距离,就为每个内容找到最相似的内容了。

approxSimilarityJoin需要传入2个embedding表格,会为每条左侧的embedding向量计算右侧与其在相同桶内的embedding向量之间距离,并且只保留距离在指定阈值之内的pair,这里左右都是同一张表,保留的向量距离阈值是0.5,距离存在Distance字段。

上面对输出结果进行了一定的整理,方便查看;另外,上述计算过程还是比较消耗内存的,如果OOM需要考虑调大spark任务的内存分配大小。

上述结果保留了每个Comment的相近Comment,即1个ID1可能对应多条ID2,因为它们距离都满足<=0.5。

第6步:为每条评论A仅保留TOP 3近似的评论B

对上述表格中的每条ID1做聚合统计,保留其Distance最近的3条ID2评论,将3条评论的信息拉平到列。

为了方便,我们使用视图表,直接写SQL完成处理:

可以看到,结果表格中每条评论的3条相似评论被拉平到列,以数组的形式存在。

我们也可以看到,word2vec生成的内容embedding向量很好的发现了一些有意思的相似评论:

电影《大鱼海棠》的评论:“美术四颗星星,剧情倒扣一颗”,与其他3部电影的如下3条评论相近::

  • 两颗星给舒淇,两颗星给陈坤。特技制作倒扣一颗星,故事情节怒扣一颗星!
  • 特效四分,景甜扣一分
  • 配音减半颗星,剧情减半颗星,制作加一颗星

不过突然发现collect_set好像把值乱序了,大家可以换成collect_list试试 :mrgreen:

最后

word2vec可以为任意序列中的item生成embedding,比如商品访问序列可以训练商品embedding,其之间的距离可以一定程度上体现出商品之间的相似度。此时,将用户的商品序列的embedding求平均可以作为用户的embedding,此时用户embedding和商品embedding之间也可以进行距离计算,体现出用户对商品的偏好,此思路可以用作线上推荐商品召回。

word2vec训练好的embedding可以导入到向量实时检索系统,实现在线的embedding高效距离计算,其原理类似LSH,可以看一下facebook开源的主流向量检索服务:faiss。

如果文章帮助您解决了工作难题,您可以帮我点击屏幕上的任意广告,或者赞助少量费用来支持我的持续创作,谢谢~