将语义搜索添加到 Elasticsearch

本教程系列将涵盖txtai的主要用例，这是一个 AI 驱动的语义搜索平台。该系列的每章都有相关代码，可也可以在colab 中使用。
colab 地址

在本文中，我们将采用第 2 部分中使用的相同 Hugging Face Dataset，在 Elasticsearch 中对其进行索引，并使用来自 txtai 的语义相似度函数对搜索结果进行排名。

安装txtai,datasets和Elasticsearch.

# Install txtai, datasets and elasticsearch python client
pip install txtai datasets elasticsearch

# Download and extract elasticsearch
wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.10.1-linux-x86_64.tar.gz
tar -xzf elasticsearch-7.10.1-linux-x86_64.tar.gz
chown -R daemon:daemon elasticsearch-7.10.1

import os
from subprocess import Popen, PIPE, STDOUT

# Start and wait for server
server = Popen(['elasticsearch-7.10.1/bin/elasticsearch'], stdout=PIPE, stderr=STDOUT, preexec_fn=lambda: os.setuid(1))

sleep 30

以下块将数据集加载到 Elasticsearch 中

from datasets import load_dataset

from elasticsearch import Elasticsearch, helpers

# 连接到 ES 实例
es = Elasticsearch(hosts=["http://localhost:9200"], timeout=60, retry_on_timeout=True)

# 下载HF数据集
dataset = load_dataset("ag_news", split="train")["text"][:50000]

# Elasticsearch 大容量缓冲区
buffer = []
rows = 0

for x, text in enumerate(dataset):
  # 文章记录
  article = {"_id": x, "_index": "articles", "title": text}

  # 缓存文章
  buffer.append(article)

  # 增加处理的文章数量
  rows += 1

  # 缓存每1000条数据加载处理一次
  if rows % 1000 == 0:
    helpers.bulk(es, buffer)
    buffer = []

    print("Inserted {} articles".format(rows), end="r")

if buffer:
  helpers.bulk(es, buffer)

print("Total articles inserted: {}".format(rows))

Elasticsearch 是一个基于令牌的搜索系统。查询和文档被解析为标记，并且使用评分算法计算最相关的查询文档匹配。默认评分算法是BM25。可以使用丰富的查询语法和Query DSL构建强大的查询。

以下部分针对 Elasticsearch 运行查询，查找前 5 个匹配项并返回与每个匹配项关联的相应文档。

from IPython.display import display, HTML

def table(category, query, rows):
    html = """
    
    @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');
    table {
      border-collapse: collapse;
      width: 900px;
    }
    th, td {
        border: 1px solid #9e9e9e;
        padding: 10px;
        font: 15px Oswald;
    }
    
    """

    html += "[%s] %s" % (category, query)
    for score, text in rows:
        html += "" % (score, text)
    html += "
Score
Text
%.4f
%s
"

    display(HTML(html))

def search(query, limit):
  query = {
      "size": limit,
      "query": {
          "query_string": {"query": query}
      }
  }

  results = []
  for result in es.search(index="articles", body=query)["hits"]["hits"]:
    source = result["_source"]
    results.append((min(result["_score"], 18) / 18, source["title"]))

  return results

limit = 3
query= "+yankees lose"
table("Elasticsearch", query, search(query, limit))

下面的代码创建了一个 Similarity 实例并定义了一个排名函数，以根据计算出的相似度对搜索结果进行排序。
ranksearch 查询 Elasticsearch 以获得更大的结果集，使用相似性实例对结果进行排名并返回前 n 个结果。

from txtai.pipeline import Similarity

def ranksearch(query, limit):
  results = [text for _, text in search(query, limit * 10)]
  return [(score, results[x]) for x, score in similarity(query, results)][:limit]

# Create similarity instance for re-ranking
similarity = Similarity("valhalla/distilbart-mnli-12-3")

现在让我们重新运行之前的搜索。

# Run the search
table("Elasticsearch + txtai", query, ranksearch(query, limit))

上面的结果在查找语义上与查询含义相似的结果方面做得更好。它不只是查找与yankees和 的匹配项lose，而是查找yankees lose.

这种组合是有效和强大的。它利用了 Elasticsearch 的高性能，同时增加了语义搜索能力。我们可能已经拥有一个大型 Elasticsearch 集群，其中包含 TB（或 PB）+ 的数据和多年的工程投资，可以解决大多数用例。对搜索结果进行语义排名是一种实用的方法。

现在来看更多比较 Elasticsearch 与 Elasticsearch + txtai 结果的示例。

for query in ["good news +economy", "bad news +economy"]:
  table("Elasticsearch", query, search(query, limit))
  table("Elasticsearch + txtai", query, ranksearch(query, limit))

https://dev.to/neuml/tutorial-series-on-txtai-ibg