Hugging Face 데이터셋을 Qdrant에 불러오기

Hugging Face 데이터셋을 Qdrant에 불러오기 (tutorials-basics-huggingface-datasets)

Qdrant를 연습하거나 시맨틱 검색 기반 애플리케이션을 만들려면 실전용 데이터가 필요해요. Hugging Face에는 ML 모델과 데이터셋을 공유하고 사용하는 플랫폼이 마련돼 있는데, Qdrant도 여기에 임베딩이 포함된 데이터셋을 함께 공개하고 있어요. 이 문서에서는 그 데이터셋을 Qdrant에 불러오는 방법을 설명해 드릴게요.

출처: Qdrant 공식문서

Hugging Face는 ML 모델과 데이터셋을 공유·사용하는 플랫폼을 제공해요. Qdrant도 임베딩과 함께 데이터셋을 공개하고 있어서, Qdrant로 연습하고 시맨틱 검색 기반 애플리케이션을 만드는 데 쓸 수 있어요. 특정 데이터셋을 보고 싶다면 알려주세요!

arxiv-titles-instructorxl-embeddings

이 데이터셋은 논문 제목으로만 생성된 임베딩을 담고 있어요. 각 벡터는 그 생성에 사용된 제목과 DOI(Digital Object Identifier)를 담은 페이로드를 갖고 있어요.

{
    "title": "Nash Social Welfare for Indivisible Items under Separable, Piecewise-Linear Concave Utilities",
    "DOI": "1612.05191"
}

데이터셋에 대한 자세한 설명은 Practice Datasets 섹션에서 찾을 수 있어요. Qdrant 스냅샷으로 데이터셋을 불러오는 쪽을 선호한다면 거기에도 링크가 있어요.

데이터셋을 불러오는 건 datasets 라이브러리의 load_dataset 함수를 쓰는 것만큼 간단해요.

from datasets import load_dataset

dataset = load_dataset("Qdrant/arxiv-titles-instructorxl-embeddings")

이 데이터셋은 2,250,000개의 벡터를 담고 있어요. 데이터셋의 피처 목록을 확인하는 방법은 다음과 같아요.

dataset.features

데이터셋 스트리밍하기 (Streaming the dataset)

데이터셋 스트리밍(dataset streaming)을 이용하면 다운로드하지 않고 데이터셋을 다룰 수 있어요. 데이터는 데이터셋을 반복(iterate)할 때 스트리밍으로 전달돼요. 자세한 내용은 Hugging Face 문서에서 읽어볼 수 있어요.

from datasets import load_dataset

dataset = load_dataset(
    "Qdrant/arxiv-titles-instructorxl-embeddings", split="train", streaming=True
)

데이터셋을 Qdrant에 불러오기

Python SDK를 사용해 데이터셋을 Qdrant에 불러올 수 있어요. 임베딩은 이미 계산되어 있으므로, 잠시 후 만들 컬렉션에 그대로 저장하면 돼요.

from qdrant_client import QdrantClient, models

client = QdrantClient("http://localhost:6333")

client.create_collection(
    collection_name="arxiv-titles-instructorxl-embeddings",
    vectors_config=models.VectorParams(
        size=768,
        distance=models.Distance.COSINE,
    ),
)

큰 데이터셋을 불러올 때는 배칭(batching)을 쓰는 게 항상 좋은 생각이에요. 그렇게 해볼게요. 데이터셋을 배치로 나누는 헬퍼 함수가 필요해요.

from itertools import islice

def batched(iterable, n):
    iterator = iter(iterable)
    while batch := list(islice(iterator, n)):
        yield batch

Python 3.12 이상을 쓰고 있다면 itertools 패키지의 batched 함수를 대신 사용할 수 있어요.

Python 버전과 상관없이 upsert 메서드로 데이터셋을 배치 단위로 Qdrant에 불러올 수 있어요.

batch_size = 100

for batch in batched(dataset, batch_size):
    ids = [point.pop("id") for point in batch]
    vectors = [point.pop("vector") for point in batch]

    client.upsert(
        collection_name="arxiv-titles-instructorxl-embeddings",
        points=models.Batch(
            ids=ids,
            vectors=vectors,
            payloads=batch,
        ),
    )

이제 컬렉션을 검색에 사용할 준비가 됐어요! Hugging Face 허브에 더 많은 데이터셋이 공개되길 원한다면 Discord로 알려주세요.