使用 Hugging Face 数据集构建 Embeddings 索引

使用 Hugging Face 数据集构建 Embeddings 索引

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

安装依赖

pip install txtai
pip install datasets

在此示例中，我们将加载ag_news数据集，该数据集是新闻文章标题的集合。这只需要一行代码！

接下来，txtai 将索引数据集的前 10,000 行。在 msmarco 上训练的模型用于计算句子嵌入。句子转换器有许多可以交换的预训练模型。

除了嵌入索引，我们还将创建一个 Similarity 实例来重新排列搜索命中的相关性。

from datasets import load_dataset

from txtai.embeddings import Embeddings
from txtai.pipeline import Similarity

def stream(dataset, field, limit):
  index = 0
  for row in dataset:
    yield (index, row[field], None)
    index += 1

    if index >= limit:
      break

def search(query):
  return [(score, dataset[uid]["text"]) for uid, score in embeddings.search(query, limit=50)]

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

# Load HF dataset
dataset = load_dataset("ag_news", split="train")

# Create embeddings model, backed by sentence-transformers & transformers
embeddings = Embeddings({"path": "sentence-transformers/msmarco-distilbert-base-v4"})
embeddings.index(stream(dataset, "text", 10000))

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

搜索数据集

现在索引已准备就绪，让我们搜索数据！以下部分运行一系列查询并显示结果。与基本的搜索引擎一样，txtai 查找标记匹配。但 txtai 的真正威力在于找到语义相似的结果。

Sentence-transformers 对信息检索有一个很好的概述，非常值得一读。

from IPython.core.display import display, HTML

def table(query, rows):
    html = """
    <style type='text/css'>
    @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;
    }
    </style>
    """

    html += "<h3>%s</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead>" % (query)
    for score, text in rows:
        html += "<tr><td>%.4f</td><td>%s</td></tr>" % (score, text)
    html += "</table>"

    display(HTML(html))

for query in ["Positive Apple reports", "Negative Apple reports", "Best planets to explore for life", "LA Dodgers good news", "LA Dodgers bad news"]:
  table(query, ranksearch(query)[:2])

参考

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