搜尋索引

FAISS 和 Elasticsearch 支援在資料集中進行搜尋。當您想要從資料集中檢索與 NLP 任務相關的特定範例時，這非常有用。例如，如果您正在進行開放領域問答（Open Domain Question Answering）任務，您可能只想返回與回答問題相關的範例。

本指南將向您展示如何為您的資料集建立索引，以便進行搜尋。

FAISS

FAISS 根據向量表示的相似度來檢索文件。在本範例中，您將使用 DPR 模型產生向量表示。

從 🤗 Transformers 下載 DPR 模型

>>> from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
>>> import torch
>>> torch.set_grad_enabled(False)
>>> ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
>>> ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

載入您的資料集並計算向量表示

>>> from datasets import load_dataset
>>> ds = load_dataset('community-datasets/crime_and_punish', split='train[:100]')
>>> ds_with_embeddings = ds.map(lambda example: {'embeddings': ctx_encoder(**ctx_tokenizer(example["line"], return_tensors="pt"))[0][0].numpy()})

使用 Dataset.add_faiss_index() 建立索引

>>> ds_with_embeddings.add_faiss_index(column='embeddings')

現在您可以使用 embeddings 索引來查詢您的資料集。載入 DPR 問題編碼器 (Question Encoder)，並使用 Dataset.get_nearest_examples() 搜尋問題。

>>> from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer
>>> q_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
>>> q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")

>>> question = "Is it serious ?"
>>> question_embedding = q_encoder(**q_tokenizer(question, return_tensors="pt"))[0][0].numpy()
>>> scores, retrieved_examples = ds_with_embeddings.get_nearest_examples('embeddings', question_embedding, k=10)
>>> retrieved_examples["line"][0]
'_that_ serious? It is not serious at all. It’s simply a fantasy to amuse\r\n'

您可以透過 Dataset.get_index() 存取索引並將其用於特殊操作，例如使用 range_search 進行查詢。

>>> faiss_index = ds_with_embeddings.get_index('embeddings').faiss_index
>>> limits, distances, indices = faiss_index.range_search(x=question_embedding.reshape(1, -1), thresh=0.95)

當您完成查詢後，請使用 Dataset.save_faiss_index() 將索引儲存到磁碟。

>>> ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')

稍後可以使用 Dataset.load_faiss_index() 重新載入它。

>>> ds = load_dataset('community-datasets/crime_and_punish', split='train[:100]')
>>> ds.load_faiss_index('embeddings', 'my_index.faiss')

Elasticsearch

與 FAISS 不同，Elasticsearch 是根據精確匹配來檢索文件。

在您的機器上啟動 Elasticsearch；如果您尚未安裝，請參閱 Elasticsearch 安裝指南。

載入您想要索引的資料集

>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')

使用 Dataset.add_elasticsearch_index() 建立索引

>>> squad.add_elasticsearch_index("context", host="localhost", port="9200")

然後，您可以使用 Dataset.get_nearest_examples() 查詢 context 索引。

>>> query = "machine"
>>> scores, retrieved_examples = squad.get_nearest_examples("context", query, k=10)
>>> retrieved_examples["title"][0]
'Computational_complexity_theory'

如果您想要重複使用該索引，請在建立索引時定義 es_index_name 參數。

>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')
>>> squad.add_elasticsearch_index("context", host="localhost", port="9200", es_index_name="hf_squad_val_context")
>>> squad.get_index("context").es_index_name
hf_squad_val_context

之後呼叫 Dataset.load_elasticsearch_index() 並提供該索引名稱，即可重新載入。

>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')
>>> squad.load_elasticsearch_index("context", host="localhost", port="9200", es_index_name="hf_squad_val_context")
>>> query = "machine"
>>> scores, retrieved_examples = squad.get_nearest_examples("context", query, k=10)

若要進行更進階的 Elasticsearch 使用，您可以透過自訂設定來指定自己的配置。

>>> import elasticsearch as es
>>> import elasticsearch.helpers
>>> from elasticsearch import Elasticsearch
>>> es_client = Elasticsearch([{"host": "localhost", "port": "9200"}])  # default client
>>> es_config = {
...     "settings": {
...         "number_of_shards": 1,
...         "analysis": {"analyzer": {"stop_standard": {"type": "standard", " stopwords": "_english_"}}},
...     },
...     "mappings": {"properties": {"text": {"type": "text", "analyzer": "standard", "similarity": "BM25"}}},
... }  # default config
>>> es_index_name = "hf_squad_context"  # name of the index in Elasticsearch
>>> squad.add_elasticsearch_index("context", es_client=es_client, es_config=es_config, es_index_name=es_index_name)

在 GitHub 上更新

資料集 (Datasets)

搜尋索引

FAISS

Elasticsearch