Indeks pencarian

FAISS dan Elasticsearch memungkinkan pencarian contoh dalam dataset. Hal ini berguna ketika Anda ingin mengambil contoh tertentu dari dataset yang relevan dengan tugas NLP Anda. Misalnya, jika Anda mengerjakan tugas Open Domain Question Answering, Anda mungkin hanya ingin mengembalikan contoh yang relevan untuk menjawab pertanyaan Anda.

Panduan ini akan menunjukkan cara membangun indeks untuk dataset Anda yang memungkinkan Anda untuk mencarinya.

FAISS

FAISS mengambil dokumen berdasarkan kemiripan representasi vektornya. Dalam contoh ini, Anda akan menghasilkan representasi vektor dengan model DPR.

Unduh model DPR dari 🤗 Transformers:

>>> 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")

Muat dataset Anda dan hitung representasi vektor:

>>> 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()})

Buat indeks dengan Dataset.add_faiss_index():

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

Sekarang Anda dapat melakukan kueri pada dataset Anda dengan indeks embeddings. Muat DPR Question Encoder, dan cari sebuah pertanyaan dengan 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'

Anda dapat mengakses indeks dengan Dataset.get_index() dan menggunakannya untuk operasi khusus, misalnya melakukan kueri menggunakan 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)

Setelah Anda selesai melakukan kueri, simpan indeks di disk dengan Dataset.save_faiss_index():

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

Muat ulang pada waktu kemudian dengan Dataset.load_faiss_index():

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

Elasticsearch

Berbeda dengan FAISS, Elasticsearch mengambil dokumen berdasarkan kecocokan yang persis.

Jalankan Elasticsearch di mesin Anda, atau lihat panduan instalasi Elasticsearch jika Anda belum menginstalnya.

Muat dataset yang ingin Anda indekskan:

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

Bangun indeks dengan Dataset.add_elasticsearch_index():

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

Lalu Anda dapat menanyakan indeks context dengan Dataset.get_nearest_examples():

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

Jika Anda ingin menggunakan ulang indeks, definisikan parameter es_index_name saat Anda membangun indeks:

>>> 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

Muat ulang nanti dengan nama indeks saat Anda memanggil 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)

Untuk penggunaan Elasticsearch yang lebih lanjut, Anda dapat menentukan konfigurasi Anda sendiri dengan pengaturan khusus:

>>> 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)

Update on GitHub

Datasets

Indeks pencarian

FAISS

Elasticsearch