VLLMRanker

VLLMRanker

vLLM으로 서빙된 reranker 모델을 사용해 문서를 쿼리와의 유사도에 따라 다시 정렬(랭킹)해 주는 컴포넌트예요.

파이프라인에서 가장 흔한 위치: 쿼리 파이프라인에서 Retriever처럼 문서 리스트를 반환하는 컴포넌트 다음 필수 init 변수: model — vLLM이 서빙하는 reranker 모델 이름 필수 run 변수: query — 쿼리 문자열 / documents — 문서 객체 리스트 출력 변수: documents — 문서 객체 리스트 / meta — 사용된 모델과 사용 정보를 담은 딕셔너리 API 레퍼런스: vLLM GitHub 링크: https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/vllm 패키지 이름: vllm-haystack

출처: 문서

본문

개요 (Overview)

vLLM은 LLM을 위한 고처리량·메모리 효율적 추론/서빙 엔진이에요. HTTP 서버를 노출하는데, VLLMRanker가 /rerank 엔드포인트를 통해 문서를 다시 랭킹하는 데 사용해요.

VLLMRanker는 api_base_url 파라미터(기본 http://localhost:8000/v1)에서 접근 가능한 vLLM 서버가 실행 중이어야 해요. 쿼리 파이프라인에서 Retriever 다음에 이 컴포넌트를 쓰면 검색된 문서를 쿼리와의 관련성에 따라 재정렬해 줘요.

top_k 파라미터로 반환할 최대 문서 수를, score_threshold 파라미터로 주어진 값 아래의 관련성 점수를 가진 문서를 제외할 수 있어요.

vLLM 서버를 --api-key로 시작했다면, Haystack의 Secret API를 통해 VLLM_API_KEY 환경 변수나 api_key init 파라미터로 API 키를 제공하세요.

호환 모델 (Compatible models)

vLLM은 다양한 reranker 모델을 지원해요. 지원되는 아키텍처와 모델 목록은 vLLM 지원 모델 문서를 참고하세요.

vLLM 특유 파라미터 (vLLM-specific parameters)

extra_parameters 딕셔너리로 vLLM 특유 파라미터를 전달할 수 있어요. 이 값들은 /rerank 엔드포인트로 보내지는 요청 본문에 병합돼요. 표준 rerank API에 없는 파라미터(예: truncate_prompt_tokens)를 넘길 때 사용해요. 자세한 내용은 vLLM rerank API 문서를 참고하세요.

ranker = VLLMRanker(
    model="BAAI/bge-reranker-base",
    extra_parameters={"truncate_prompt_tokens": 256},
)

메타 필드 임베딩 (Embedding meta fields)

재랭킹할 때 문서 내용과 함께 제목 같은 메타 정보를 포함하면 이점이 있는 유스케이스가 있어요. 포함할 메타 필드 이름을 meta_fields_to_embed 파라미터로 전달하면 돼요. 이 값들은 meta_data_separator를 사용해 문서 내용과 연결돼요.

ranker = VLLMRanker(
    model="BAAI/bge-reranker-base",
    meta_fields_to_embed=["title"],
    meta_data_separator="\n",
)

사용법 (Usage)

VLLMRanker를 쓰려면 vllm-haystack 패키지를 설치해요:

pip install vllm-haystack

vLLM 서버 시작 (Starting the vLLM server)

이 컴포넌트를 쓰기 전에 reranker 모델로 vLLM 서버를 시작하세요:

vllm serve BAAI/bge-reranker-base

서버 옵션에 대한 자세한 내용은 vLLM CLI 문서를 참고하세요.

단독으로 사용하기 (On its own)

from haystack import Document
from haystack_integrations.components.rankers.vllm import VLLMRanker

ranker = VLLMRanker(model="BAAI/bge-reranker-base")
docs = [
    Document(content="The capital of Brazil is Brasilia."),
    Document(content="The capital of France is Paris."),
]
result = ranker.run(query="What is the capital of France?", documents=docs)
print(result["documents"][0].content)
# The capital of France is Paris.

파이프라인 안에서 사용하기 (In a pipeline)

from haystack import Document, Pipeline
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.rankers.vllm import VLLMRanker

docs = [
    Document(content="Paris is in France"),
    Document(content="Berlin is in Germany"),
    Document(content="Lyon is in France"),
]
document_store = InMemoryDocumentStore()
document_store.write_documents(docs)

retriever = InMemoryBM25Retriever(document_store=document_store)
ranker = VLLMRanker(model="BAAI/bge-reranker-base")

document_ranker_pipeline = Pipeline()
document_ranker_pipeline.add_component(instance=retriever, name="retriever")
document_ranker_pipeline.add_component(instance=ranker, name="ranker")
document_ranker_pipeline.connect("retriever.documents", "ranker.documents")

query = "Cities in France"
result = document_ranker_pipeline.run(
    data={
        "retriever": {"query": query, "top_k": 3},
        "ranker": {"query": query, "top_k": 2},
    },
)
print(result["ranker"]["documents"][0])
# Document(id=..., content: 'Paris is in France', score: ...)

더 알아보기 (Learn more)