임베딩 모델을 활용한 재순위화

RAG의 재순위화를 위한 Python 코드

Page content

리랭킹 은 검색 증강 생성 (RAG) 시스템 에서 검색 (Retrieving) 과 생성 (Generating) 사이에 위치하는 두 번째 단계입니다.

디지털 공간의 전기 큐브

위 이미지는 Flux-1 dev 가 Electric cubes in digital space를 어떻게 상상하는지를 보여줍니다.

RAG 시스템 구축에 대한 전체 가이드는 검색 증강 생성 (RAG) 튜토리얼: 아키텍처, 구현 및 프로덕션 가이드 를 참고하세요.

이 글은 임베딩 모델을 활용한 리랭킹에 대해 다룹니다. 다른 접근 방식으로는 Ollama 와 Qwen3 임베딩 LLM 을 활용한 텍스트 리랭킹 - Go 언어 및 Ollama 와 Qwen3 리랭커 모델을 활용한 문서 리랭킹 - Go 언어 을 확인해 보실 수 있습니다.

리랭킹을 통한 검색

초기부터 문서를 임베딩 형태로 벡터 DB 에 저장한다면, 검색을 수행하면 즉시 유사한 문서 목록을 얻을 수 있습니다.

독립형 리랭킹 (Standalone reranking)

하지만 인터넷에서 문서를 먼저 다운로드한 경우, 검색 시스템의 응답은 검색 제공자의 선호도/알고리즘, 스폰서 콘텐츠, SEO 최적화 등에 영향을 받을 수 있으므로 검색 후 리랭킹이 필요합니다.

저는 다음과 같은 작업을 수행했습니다.

  • 검색 쿼리에 대한 임베딩을 얻습니다.
  • 각 문서에 대한 임베딩을 얻습니다. (문서는 어차피 8 토큰을 초과하지 않을 것으로 예상되었습니다)
  • 쿼리와 각 문서 임베딩 간의 유사도를 계산합니다.
  • 이 유사도를 기준으로 문서를 정렬합니다.

여기에는 벡터 DB 가 필요 없으므로 바로 시작해 보겠습니다.

샘플 코드

Ollama 에 연결하고 langchain 의 cosine_similarity 함수를 사용하기 위해 Langchain 을 사용합니다. 유사도 측정 기준으로 필터링할 수 있지만, 도메인과 임베딩 LLM 에 따라 임계값이 달라질 수 있음을 유의하세요.

이 코드가 여러분께 조금이라도 도움이 되기를 바랍니다. 복사/붙여넣기/원하는 대로 사용 가능한 라이선스입니다. 감사합니다.

from langchain_core.documents import Document
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.utils.math import cosine_similarity
import numpy as np


def cosine_distance(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return 1.0 - cosine_similarity(a, b)

def compute_score(vectors: np.ndarray) -> float:
    score = cosine_distance(vectors[0].reshape(1, -1), vectors[1].reshape(1, -1)).item()
    return score

def list_to_array(lst):
    return np.array(lst, dtype=float)   

def compute_scorel(lists) -> float:
    v1 = list_to_array(lists[0])
    v2 = list_to_array(lists[1])
    return compute_score([v1, v2])

def filter_docs(emb_model_name, docs, query, num_docs):
    content_arr = [doc.page_content for doc in docs]

    ollama_emb = OllamaEmbeddings(
        model=emb_model_name
    )

    docs_embs = ollama_emb.embed_documents(content_arr)
    query_embs = ollama_emb.embed_query(query)
    sims = []
    for i, emb in enumerate(docs_embs):
        idx = docs[i].id
        s = compute_scorel([query_embs, docs_embs[i]])
        simstr = str(round(s, 4))
        docs[i].metadata["sim"] = simstr
        sim = {
            "idx": idx,
            "i": i,
            "sim": s,
        }
        sims.append(sim)

    sims.sort(key=sortFn)

    sorted_docs = [docs[x["i"]] for x in sims]
    filtered_docs = sorted_docs[:num_docs]
    return filtered_docs

최적의 임베딩 모델

제 작업에 현재 가장 좋은 임베딩 모델은 bge-large:335m-en-v1.5-fp16 입니다.

다음으로는 nomic-embed-text:137m-v1.5-fp16 과 jina/jina-embeddings-v2-base-en:latest 가 2 위를 차지했습니다.

하지만 본인의 도메인과 쿼리에 맞춰 직접 테스트해 보시는 것이 좋습니다.

유용한 링크