MariaDBEmbeddingRetriever

MariaDBEmbeddingRetriever

MariaDB Document Store와 호환되는 임베딩 기반 Retriever예요.

출처: MariaDBEmbeddingRetriever

본문

개요

MariaDBEmbeddingRetriever는 MariaDBDocumentStore와 호환되는 임베딩 기반 Retriever예요. 쿼리와 문서 임베딩을 비교하고 MariaDB의 네이티브 MHNSW 벡터 인덱스를 사용해 쿼리와 가장 관련성 높은 문서를 가져와요.

파이프라인에서 MariaDBEmbeddingRetriever를 쓸 때는 임베딩이 준비되어 있어야 해요. 인덱싱 파이프라인에 Document Embedder를, 쿼리 파이프라인에 Text Embedder를 추가하면 돼요.

query_embedding 외에도 Retriever는 top_k(가져올 최대 문서 수), 검색 공간을 좁히는 filters 같은 선택적 파라미터를 받아요.

벡터 인덱스

빠른 근사 최근접 이웃 검색을 위해선 MariaDBDocumentStore를 create_vector_index=True로 초기화해야 해요. 이렇게 하면 테이블 생성 때 MHNSW 인덱스가 만들어지지만, 모든 문서에 null이 아닌 임베딩이 필요해요. embedding_dimension과 distance 파라미터도 테이블 생성 때만(또는 recreate_table=True와 함께) 적용돼요.

설치

MariaDB 11.7 인스턴스를 빠르게 설정하려면 Docker를 사용하면 돼요.

docker run -d -p 3306:3306 \
 -e MARIADB_ROOT_PASSWORD=secret \
 -e MARIADB_DATABASE=haystack \
 -e MARIADB_USER=haystack \
 -e MARIADB_PASSWORD=secret \
 mariadb:11.7

시스템 라이브러리와 통합 패키지를 설치해요.

# Ubuntu / Debian
sudo apt-get install -y libmariadb-dev

pip install mariadb-haystack

아래 파이프라인 예시는 Sentence Transformers 임베더도 사용해요.

pip install sentence-transformers-haystack

사용법

단독 사용:

import os
from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore
from haystack_integrations.components.retrievers.mariadb import (
    MariaDBEmbeddingRetriever,
)

os.environ["MARIADB_USER"] = "haystack"
os.environ["MARIADB_PASSWORD"] = "secret"

document_store = MariaDBDocumentStore(embedding_dimension=768)
retriever = MariaDBEmbeddingRetriever(document_store=document_store)

# using a fake vector to keep the example simple
retriever.run(query_embedding=[0.1] * 768)

파이프라인 안에서:

import os
from haystack import Document, Pipeline
from haystack_integrations.components.embedders.sentence_transformers import (
    SentenceTransformersTextEmbedder,
    SentenceTransformersDocumentEmbedder,
)
from haystack.document_stores.types import DuplicatePolicy

from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore
from haystack_integrations.components.retrievers.mariadb import (
    MariaDBEmbeddingRetriever,
)

os.environ["MARIADB_USER"] = "haystack"
os.environ["MARIADB_PASSWORD"] = "secret"

document_store = MariaDBDocumentStore(
    embedding_dimension=768,
    distance="cosine",
)

documents = [
    Document(content="There are over 7,000 languages spoken around the world today."),
    Document(
        content="Elephants have been observed to recognize themselves in mirrors."
    ),
    Document(
        content="Bioluminescent waves can be seen in the Maldives and Puerto Rico."
    ),
]

document_embedder = SentenceTransformersDocumentEmbedder()
documents_with_embeddings = document_embedder.run(documents)

document_store.write_documents(
    documents_with_embeddings.get("documents"),
    policy=DuplicatePolicy.OVERWRITE,
)

query_pipeline = Pipeline()
query_pipeline.add_component("text_embedder", SentenceTransformersTextEmbedder())
query_pipeline.add_component(
    "retriever",
    MariaDBEmbeddingRetriever(document_store=document_store),
)
query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")

result = query_pipeline.run(
    {"text_embedder": {"text": "How many languages are there?"}}
)
print(result["retriever"]["documents"][0])

더 알아보기 (Learn more)