#
허깅페이스 허브를 백엔드로 사용한 벡터 검색

작성자: Martin Elstner , 번역: 최용빈

허깅페이스의 데이터셋은 파켓(parquet) 파일에 의존합니다. 우리는 빠른 인메모리 데이터베이스 시스템인 DuckDB를 사용하여 이 파일들과 상호작용할 수 있습니다. DuckDB의 기능 중 하나는 벡터 유사도 검색으로, 인덱스 유무에 관계없이 사용할 수 있습니다.

##
의존성 설치

    !pip install datasets duckdb sentence-transformers model2vec -q

##
데이터셋에 대한 임베딩 생성

먼저, 검색할 데이터셋에 대한 임베딩을 생성해야 합니다. 우리는 `sentence-transformers` 라이브러리를 사용하여 데이터셋에 대한 임베딩을 생성할 것입니다.

    from sentence_transformers import SentenceTransformer
    from sentence_transformers.models import StaticEmbedding

    static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-8M")
    model = SentenceTransformer(modules=[static_embedding])

이제 허브에서 ai-blueprint/fineweb-bbc-news 데이터셋을 로드해 보겠습니다.

    from datasets import load_dataset

    ds = load_dataset("ai-blueprint/fineweb-bbc-news")

이제 데이터셋에 대한 임베딩을 생성할 수 있습니다. 일반적으로 정밀도(precision)를 잃지 않기 위해 데이터를 더 작은 배치로 나누고 싶을 수 있지만, 이 예제에서는 데이터셋의 전체 텍스트에 대한 임베딩만 생성하겠습니다.

    def create_embeddings(batch):
        embeddings = model.encode(batch["text"], convert_to_numpy=True)
        batch["embeddings"] = embeddings.tolist()
        return batch

    ds = ds.map(create_embeddings, batched=True)

이제 임베딩이 포함된 데이터셋을 허브에 다시 업로드할 수 있습니다.

    ds.push_to_hub("ai-blueprint/fineweb-bbc-news-embeddings")

##
허깅페이스 허브에서 벡터 검색

이제 `duckdb`를 사용하여 데이터셋에서 벡터 검색을 수행할 수 있습니다. 이때 인덱스를 사용하거나 사용하지 않을 수 있습니다. 인덱스를 활용하지 않고 검색하는 것은 더 느리지만 더 정확하고, 인덱스를 활용하여 검색하는 것은 더 빠르지만 덜 정확합니다.

###
인덱스를 활용하지 않고 검색하기

인덱스를 활용하지 않고 검색하려면 `duckdb` 라이브러리를 사용하여 데이터셋에 연결하고 벡터 검색을 수행할 수 있습니다. 이는 느린 작업지만 일반적으로 약 10만 행까지의 작은 데이터셋에서는 충분히 빠르게 동작합니다. 즉, 우리가 사용하는 데이터셋에 대해서 쿼리를 날리는 것은 다소 느릴 것입니다.

    import duckdb
    from typing import List

    def similarity_search_without_duckdb_index(
        query: str,
        k: int = 5,
        dataset_name: str = "ai-blueprint/fineweb-bbc-news-embeddings",
        embedding_column: str = "embeddings",
    ):
        # 인덱스를 위해 사용한 모델과 동일한 모델을 사용
        query_vector = model.encode(query)
        embedding_dim = model.get_sentence_embedding_dimension()

        sql = f"""
            SELECT
                *,
                array_cosine_distance(
                    {embedding_column}::float[{embedding_dim}],
                    {query_vector.tolist()}::float[{embedding_dim}]
                ) as distance
            FROM 'hf://datasets/{dataset_name}/**/*.parquet'
            ORDER BY distance
            LIMIT {k}
        """
        return duckdb.sql(sql).to_df()

    similarity_search_without_duckdb_index("What is the future of AI?")

###
인덱스를 활용하여 검색하기

이 접근법은 데이터셋의 로컬 복사본을 생성하고 이를 사용하여 인덱스를 생성합니다. 약간의 경미한 오버헤드가 있지만 한번 생성한 후에는 검색 속도가 상당히 향상될 것입니다.

    import duckdb

    def _setup_vss():
        duckdb.sql(
            query="""
            INSTALL vss;
            LOAD vss;
            """
        )
    def _drop_table(table_name):
        duckdb.sql(
            query=f"""
            DROP TABLE IF EXISTS {table_name};
            """
        )

    def _create_table(dataset_name, table_name, embedding_column):
        duckdb.sql(
            query=f"""
            CREATE TABLE {table_name} AS
            SELECT *, {embedding_column}::float[{model.get_sentence_embedding_dimension()}] as {embedding_column}_float
            FROM 'hf://datasets/{dataset_name}/**/*.parquet';
            """
        )

    def _create_index(table_name, embedding_column):
        duckdb.sql(
            query=f"""
            CREATE INDEX my_hnsw_index ON {table_name} USING HNSW ({embedding_column}_float) WITH (metric = 'cosine');
            """
        )

    def create_index(dataset_name, table_name, embedding_column):
        _setup_vss()
        _drop_table(table_name)
        _create_table(dataset_name, table_name, embedding_column)
        _create_index(table_name, embedding_column)

    create_index(
        dataset_name="ai-blueprint/fineweb-bbc-news-embeddings",
        table_name="fineweb_bbc_news_embeddings",
        embedding_column="embeddings"
    )

이제 인덱스를 사용하여 벡터 검색을 수행할 수 있으며, 결과는 즉시 반환됩니다.

    def similarity_search_with_duckdb_index(
        query: str,
        k: int = 5,
        table_name: str = "fineweb_bbc_news_embeddings",
        embedding_column: str = "embeddings"
    ):
        embedding = model.encode(query).tolist()
        return duckdb.sql(
            query=f"""
            SELECT *, array_cosine_distance({embedding_column}_float, {embedding}::FLOAT[{model.get_sentence_embedding_dimension()}]) as distance
            FROM {table_name}
            ORDER BY distance
            LIMIT {k};
        """
        ).to_df()

    similarity_search_with_duckdb_index("What is the future of AI?")

이 쿼리는 응답시간을 30초에서 1초 미만으로 줄이며, 무거운 벡터 검색 엔진을 배포할 필요가 없고 저장소는 허브에서 처리됩니다.

##
결론

우리는 `duckbd`를 사용하여 허브에서 벡터 검색을 수행하는 방법을 보았습니다. 10만 행 미만의 작은 데이터셋의 경우, 허브를 벡터 검색 벡엔드로 사용하여 인덱스 없이 벡터 검색을 수행할 수 있지만, 더 큰 데이터셋의 경우 `vss` 확장 프로그램으로 인덱스를 생성하면서 로컬 검색을 수행하고 허브를 스토리지 백엔드로 사용해야 합니다.

##
더 알아보기

- 허깅페이스에서의 벡터 검색
- DuckDB를 사용한 벡터 검색 인덱싱

Update on GitHub