Datasets documentation
검색 인덱스
검색 인덱스
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)