使用 Elasticsearch 计算余弦相似度

我是M2的marutaku。我的日历一直都很空，感到很困扰。请大家尽快报名参加活动。这次我想和大家分享一下我以前在尝试建立图像搜索系统时，使用Elasticsearch计算余弦相似度的故事。

关于余弦相似度

余弦相似度用于比较向量之间的相似度。由于公式已在其他网站上经常出现，故省略。在使用Word2Vec时，用于衡量单词的相似度。若使用gensim时，可以使用gensim.models.Word2Vec.most_similar来计算余弦相似度。然而，用Python编写余弦相似度比较的处理可能会变得较慢。
这时候就需要引入Elasticsearch。

关于Elasticsearch

Elasticsearch是一个开源的搜索引擎，可以快速地从大量数据中进行搜索。它被用于全文搜索、日志收集和可视化等方面。还提供了Docker环境，可以轻松地尝试使用，请务必试试看。

从Elasticsearch7.0开始，如上所述，可以存储向量，并且从Elasticsearch7.3开始，开始支持使用向量进行文档评分。这使得在Elasticsearch中可以计算余弦相似度。

准备Elasticsearch。

由于Elasticsearch已经有Docker映像可用，所以这次我想使用该映像。以下是用于docker-compose的docker-compose.yml文件的内容。

version: "2.3"
services:
    elasticsearch:
        image: "elasticsearch:7.5.1"
        ports:
            - "9200:9200"
        volumes:
            - es-data:/usr/share/elasticsearch/data
        tty: true
        environment:
            discovery.type: single-node
volumes:
    es-data:
        driver: local

在Elasticsearch中，你需要像在MySQL中定义表格一样，事先定义好数据的结构。这个定义被称为Mapping。以下是我在过去创建图像搜索系统时使用的Mapping。

向量需要使用dense_vector字段类型来定义，并且需要事先指定向量的长度。

{
    "mappings": {
        "properties": {
            "item_url": {
                "type": "text"
            },
            "item_description": {
                "type": "text"
            },
            "image_vector": {
                "type": "dense_vector",
                "dims": 512
            }
        }
    }
}

用Python插入数据

我打算使用Python来尝试输入数据。
由于提供了一个用于在Python中操作Elasticsearch的模块，所以首先我们需要安装该模块。

pip install elasticsearch

将数据放入Elasticsearch中。

from elasticsearch import Elasticsearch

INDEX_NAME = 'test_index'
es_client = Elasticsearch(host=es_host, port=es_port)

def insert(tensor, item_url, item_description):
    es_client.index(index=INDEX_NAME, body={
        "item_url": item_url, "item_description": item_description, "image_vector": tensor
    })

试着搜索一下

如果能够存储数据，我将尝试进行实际搜索。

def search(vector):
        query = {
            "script_score": {
                "query": {"match_all": {}},
                "script": {
                    "source": "cosineSimilarity(params.image_vector, doc['image_vector']) + 1.0",
                    "params": {"image_vector": vector}
                }
            }
        }
        response = es_client.search(
            index=INDEX_NAME,
            body={
                "size": 5,
                "query": query,
                "_source": {"includes": ["item_url", "item_description"]}
            }
        )
        return response

在我的环境中，使用这个方法可以在大约40毫秒的时间内从约1万个512维向量中搜索到相似的结果。我认为这个速度非常快。

结束了

这次我尝试使用Elasticsearch来计算余弦相似度。
它比使用Python更快，所以请务必试一试。