如何优雅的使用腾讯词向量: 基于redis集群&elasticsearch的姿势

  2018年10月份，腾讯AI Lab开源了大规模高质量的中文词向量数据，包含了8824331个常用词的向量表示，维度为200。当前，向量表示已经成为nlp的重要基础功能，从我个人角度而言，其地位等同于搜索引擎中的分词功能，是注入词性标注、命名实体识别、情感分类等后续任务的基础步骤。关于腾讯词向量的进一步细节大家可以参考tencent embedding。
  关于词向量的使用，腾讯AI Lab给出的官方姿势是利用gensim加载模型，然后可以使用gensim的各种API进行其他任务，词向量的提取，相似词的查找等等。示例代码为，因为gensim版本问题，与官方代码略有出路，本机gensim版本为3.5.0:

# 腾讯词向量主页示例：from gensim.models.word2vec import KeyedVectors 
from gensim.models import KeyedVectors
wv_from_text = KeyedVectors.load_word2vec_format(file, binary=False)

由于腾讯词向量较大，数据解压后得到的词表文件为16G，因此每次加载都要十分钟左右，另外对于代码的调试也十分不方便。所以想到了用缓存的方式将词及对应向量进行存储，至于相似词等高级查找功能可以利用向量索引工具来实现(如annoy、faiss等，后续将会有文章探讨向量搜索)。
  首先想到的是利用redis键值数据库来缓存词表，由于redis是内存存储，所以博主用了三台机器搭建了redis集群，进行数据均衡，虽然采用redis集群部署可以提供较高的读写性能，该方案需要较多的内存资源。还有一种方法，即牺牲一定的读取性能，但对硬件的需求较低，即利用elasticsearch来进行缓存。下面将详细介绍两种方案。

Redis方案(需要了解Redis集群搭建方法)

  博主采用python读取本地词向量文件，然后直接写入redis集群，最终耗时1个小时左右，三个节点最终占用的内存分别为6.01G、6.02G、6.02G，词数量分别为：2942462个、2941404个、2940465个，写入耗时约为半小时(可能不准确，当时忘记记录准确时间了)。示例写入代码：

# -*- coding:utf-8 -*-
import codecs
import rediscluster


class Vector2Redis(object):
    def __init__(self):
        self.redis_client = rediscluster.StrictRedisCluster(
            startup_nodes=[
                {"host": "ip1", "port": post1},
                {"host": "ip2", "port": port2},
                {"host": "ip3", "port": port3},
            ])

    def write_to_redis(self):
        f = codecs.open("../data/Tencent_AILab_ChineseEmbedding.txt", "r", "utf-8")
        for line in f:
            try:
                split_line = line.strip().split(" ")
                word = split_line[0]
                vector = ",".join(split_line[1:])
                self.redis_client.set(word, vector)
            except Exception as exc:
                print(word, exc)


if __name__ == "__main__":
    vr = Vector2Redis()
    vr.write_to_redis()

其中rediscluster模块需要安装redis-py-cluster包。在调用时，创建一个redis集群连接，然后利用redis client的get api即可获取指定词的向量化表示，实测结果为10000个词同步获取，耗时2.5s，读取性能较好。

Elasticsearch方案(需要对es有一定了解)

  elasticsearch进行缓存不需要较多的内存资源，因为主要基于lucene索引，索引文件存于磁盘，但索引文件的大小和索引的配置有重要关系，博主经过尝试后，将mapping(类似于关系型数据的schema)最终定义为：

{
  "tencent_w2v": {
    "mappings": {
      "tencent_w2v": {
        "_all": {
          "enabled": false
        },
        "properties": {
          "vector": {
            "type": "keyword",
            "index": false
          },
          "word": {
            "type": "keyword"
          }
        }
      }
    },
    "settings": {
      "index": {
        "number_of_shards": "3",
        "number_of_replicas": "0"
      }
    }
  }
}

其中关闭了_all字段，分片数设为了3，副本数设为了0，并且由于我们只要能够实现类似于redis一样通过词将向量召回的功能即可，因此将word设定为keyword类型，并且对vector即向量字段不建立索引，这样可以有效较小最终索引体积，这种设置下最终的索引大小为27.2G。建立索引推送数据的示例代码为:

# -*- coding:utf-8 -*-
import codecs
import elasticsearch
from elasticsearch import helpers


class Vector2ES(object):
    def __init__(self):
        self.es_client = elasticsearch.Elasticsearch(
            hosts="127.0.0.1",
            port=9201
        )

    def write_to_es(self):
        f = codecs.open("../data/Tencent_AILab_ChineseEmbedding.txt", "r", "utf-8")
        actions = []
        for line in f:
            try:
                split_line = line.strip().split(" ")
                word = split_line[0]
                vector = ",".join(split_line[1:])
                action = {
                    "_index": "tencent_w2v",
                    "_type": "tencent_w2v",
                    "_source": {
                        "word": word,
                        "vector": vector
                    }
                }
                actions.append(action)
                if len(actions) == 1000:
                    helpers.bulk(self.es_client, actions, index="tencent_w2v")
                    actions = []
            except Exception as exc:
                print(exc, line)
        if len(actions) > 0:
            helpers.bulk(self.es_client, actions, index="tencent_w2v")


if __name__ == "__main__":
    vr = Vector2ES()
    vr.write_to_es()

  将词向量全部推送至elasticsearch建立索引后，调用词向量实际等价于通过搜索的方式获取，因此需要参照elasticsearch的搜索方法，示例代码为:

# -*- coding:utf-8 -*-
import jieba
import numpy
import elasticsearch


class ESSearch(object):
    def __init__(self):
        self.es_client = elasticsearch.Elasticsearch("127.0.0.1:9201")

    def encoding(self, word):
        result = self.es_client.search(index="tencent_w2v", doc_type="tencent_w2v",
                                       body={"_source": "vector", "query": {"term": {"word": word}}})
        hits = result.get("hits")
        if hits.get("hits"):
            temp_vector = hits.get("hits")[0].get("_source").get("vector")
            print(temp_vector)
            if temp_vector:
                split_temp_vector = temp_vector.split(",")
                split_temp_vector = numpy.array([float(item) for item in split_temp_vector])
        return split_temp_vector


if __name__ == "__main__":
    es_search = ESSearch()
    es_search.encoding("毛笔")

经过测试，单个词调用时间为10ms左右，同一个词调用1000次，时间约为1.7s，这是由于elasticsearch本身会做缓存，所以多次调用同一词时间消耗减小许多，但仍比redis慢8倍左右。

结论

  腾讯词向量的发布，使得大家可以方便进行各种embedding工作，但由于数据较大，因此调用及调试不是特别方便。这里给出了基于redis及elasticsearch的两种解决思路，但仅限于将原始词向量文件进行缓存，更高级的功能请继续关注博客，后续将会在此基础上，讨论向量搜索相关技术。