멀티 홉 검색 튜토리얼

멀티 홉 검색 튜토리얼 (Multi-Hop Retrieval)

여러 하위 모듈을 가진 dspy.Module을 만드는 빠른 예시를 함께 살펴볼게요. 대상 작업은 멀티 홉 검색(multi-hop search)이에요.

설치는 pip install -U dspy로 최신 DSPy를 받아서 따라오시면 되고, 추가로 pip install datasets도 필요해요.

권장사항: 내부에서 무슨 일이 벌어지는지 이해하려면 MLflow Tracing을 설정해 두세요.

출처: Multi-Hop Retrieval

MLflow DSPy 통합

MLflow는 DSPy와 기본적으로 통합되는 LLMOps 도구예요. 네 단계로 쉽게 설정할 수 있어요. MLflow를 설치하고(%pip install mlflow>=2.20), mlflow ui --port 5000으로 UI를 띄우고, mlflow.set_tracking_uri("http://localhost:5000")과 mlflow.dspy.autolog()로 연결하면 돼요. 자세한 내용은 MLflow DSPy 문서를 참고하세요.

모델 설정

여기서는 매개변수가 80억 개인 Meta의 Llama-3.1-8B-Instruct라는 작은 로컬 LM을 쓸게요. 8B 모델은 Ollama로 노트북에서, SGLang으로 GPU 서버에서, 또는 Databricks·Together 같은 호스팅 프로바이더로 운영할 수 있어요. 이 작은 모델을 메인 LM으로 설정할게요. 또 더 큰 LM인 GPT-4o를 teacher로 설정하는데, 이걸 아주 적은 횟수만 호출해 작은 LM을 가르치는 데 써요. 엄밀히는 필수는 아니에요. 작은 모델은 DSPy에서 이런 작업을 스스로 학습할 수 있지만, 큰 teacher를 쓰면 초기 시스템이나 옵티마이저 설정이 덜 중요해져서 마음이 편해져요.

import dspy

lm = dspy.LM('<your_provider>/Llama-3.1-8B-Instruct', max_tokens=3000)
gpt4o = dspy.LM('openai/gpt-4o', max_tokens=3000)

dspy.configure(lm=lm)

의존성 설치와 데이터 내려받기

검색에 가벼운 BM25S 라이브러리를 쓸게요. 이 구성요소는 원하는 것으로 바꿔도 돼요.

> pip install -U bm25s PyStemmer "jax[cpu]"

다음으로 2017년 기준 위키백과 5,000,000개 페이지의 초록(첫 문단) 스냅샷을 내려받아 검색 코퍼스로 쓸게요. 압축 상태로 500MB라 내려받고 푸는 데 2~3분 걸릴 수 있어요.

from dspy.utils import download

download("https://huggingface.co/dspy/cache/resolve/main/wiki.abstracts.2017.tar.gz")
!tar -xzvf wiki.abstracts.2017.tar.gz

코퍼스를 로드할게요.

import orjson
corpus = []

with open("wiki.abstracts.2017.jsonl") as f:
    for line in f:
        line = orjson.loads(line)
        corpus.append(f"{line['title']} | {' '.join(line['text'])}")

len(corpus)

5,233,330개 문서가 나와요. 이제 BM25 검색용으로 인덱스할게요. 2~3분 걸려요.

import bm25s
import Stemmer

stemmer = Stemmer.Stemmer("english")
corpus_tokens = bm25s.tokenize(corpus, stopwords="en", stemmer=stemmer)

retriever = bm25s.BM25(k1=0.9, b=0.4)
retriever.index(corpus_tokens)

HoVer 데이터셋 로드

HoVer 멀티 홉 작업에서 예시를 로드할게요. 입력은 (정말!) 복잡한 주장(claim)이고, 우리가 찾는 출력은 그 주장을 사실 확인(fact-check)하는 데 필요한 위키백과 페이지 집합이에요.

import random
from dspy.datasets import DataLoader

kwargs = dict(fields=("claim", "supporting_facts", "hpqa_id", "num_hops"), input_keys=("claim",))
hover = DataLoader().from_huggingface(dataset_name="vincentkoc/hover-parquet", split="train", trust_remote_code=True, **kwargs)

hpqa_ids = set()
hover = [
    dspy.Example(claim=x.claim, titles=list(set([y["key"] for y in x.supporting_facts]))).with_inputs("claim")
    for x in hover
    if x["num_hops"] == 3 and x["hpqa_id"] not in hpqa_ids and not hpqa_ids.add(x["hpqa_id"])
]

random.Random(0).shuffle(hover)
trainset, devset, testset = hover[:200], hover[200:500], hover[650:]

num_hops == 3인 예시만 골라 중복 hpqa_id를 제거하고, train 200 / dev 300 / test는 나머지로 나눠요.

이 작업의 예시를 하나 볼게요.

example = trainset[0]

print("Claim:", example.claim)
print("Pages that must be retrieved:", example.titles)
  • Claim: "이 감독은 Miss Potter 작업으로 유명하다. 그가 'Babe'로 후보에 오른 상은 영화예술과학아카데미가 수여한다."
  • 검색해야 할 페이지: ['Miss Potter', 'Chris Noonan', 'Academy Award for Best Director']

한 주장을 사실 확인하려면 여러 페이지를 잇달아 찾아야 하는 멀티 홉 작업이에요.

검색 함수와 멀티 홉 프로그램 정의

위키백과에서 검색하는 함수를 BM25 인덱스로 정의할게요.

def search(query: str, k: int) -> list[str]:
    tokens = bm25s.tokenize(query, stopwords="en", stemmer=stemmer, show_progress=False)
    results, scores = retriever.retrieve(tokens, k=k, n_threads=1, show_progress=False)
    run = {corpus[doc]: float(score) for doc, score in zip(results[0], scores[0])}
    return run

이제 DSPy에서 멀티 홉 프로그램을 정의할게요. 아주 단순해요. claim을 받아 titles: list[str] 목록을 만들어 내요. generate_query와 append_notes 두 하위 모듈로 동작해요.

class Hop(dspy.Module):
    def __init__(self, num_docs=10, num_hops=4):
        self.num_docs, self.num_hops = num_docs, num_hops
        self.generate_query = dspy.ChainOfThought('claim, notes -> query')
        self.append_notes = dspy.ChainOfThought('claim, notes, context -> new_notes: list[str], titles: list[str]')

    def forward(self, claim: str) -> list[str]:
        notes = []
        titles = []

        for _ in range(self.num_hops):
            query = self.generate_query(claim=claim, notes=notes).query
            context = search(query, k=self.num_docs)
            prediction = self.append_notes(claim=claim, notes=notes, context=context)
            notes.extend(prediction.new_notes)
            titles.extend(prediction.titles)
        
        return dspy.Prediction(notes=notes, titles=list(set(titles)))

forward는 hop마다 1) 지금까지의 notes를 바탕으로 generate_query가 검색 쿼리를 만들고, 2) 그 쿼리로 BM25 검색해 context를 얻고, 3) append_notes가 새 메모와 제목을 추가해요. 이 과정을 num_hops(4)번 반복해요.

평가 metric 설정

top5_recall metric을 정의할게요. 프로그램이 돌려준 top-5 제목 중에 정답 페이지(항상 3개)가 몇 개나 들어있는지의 비율을 돌려줘요.

def top5_recall(example, pred, trace=None):
    gold_titles = example.titles
    recall = sum(x in pred.titles[:5] for x in gold_titles) / len(gold_titles)

    # If we're "bootstrapping" for optimization, return True if and only if the recall is perfect.
    if trace is not None:
        return recall >= 1.0
    
    # If we're just doing inference, just measure the recall.
    return recall

evaluate = dspy.Evaluate(devset=devset, metric=top5_recall, num_threads=16, display_progress=True, display_table=5)

trace is not None이면 최적화용 부트스트래핑이므로 recall이 완벽할 때(≥1.0)만 참을 돌려줘요. 추론 때는 그냥 recall 값을 측정해요.

기성품 프로그램을 평가해 볼게요.

evaluate(Hop())

결과는 Average Metric: 93.99999999999993 / 300 (31.3%) 예요. dev set에서 top-5 recall이 약 31%예요. 평가 중 일부 항목은 출력 필드가 불완전해서 오류가 발생하기도 했어요. 표를 보면 예를 들어 "Finding Dory 감독이 A Bug's Life 공동 연출했다"는 주장에서 [Finding Dory, A Bug's Life]를 찾아 3개 정답 중 2개를 맞춰 recall 0.667을 내는 걸 볼 수 있어요.

MIPROv2로 최적화

Hop() 프로그램 안의 두 프롬프트를 함께 최적화해 recall을 최대화할게요. 약 35분 걸리고, Llama-3.1-8B를 최적화하기 위해 GPT-4o에 약 $5 상당의 호출을 만들 수 있어요.

models = dict(prompt_model=gpt4o, teacher_settings=dict(lm=gpt4o))
tp = dspy.MIPROv2(metric=top5_recall, auto="medium", num_threads=16, **models)

kwargs = dict(minibatch_size=40, minibatch_full_eval_steps=4)
optimized = tp.compile(Hop(), trainset=trainset, max_bootstrapped_demos=4, max_labeled_demos=4, **kwargs)

prompt_model과 teacher_settings에 모두 GPT-4o를 써서 작은 Llama가 좋은 예시를 부트스트랩하도록 도와요.

최적화 후 다시 평가할게요.

evaluate(optimized)

결과는 Average Metric: 177.33333333333334 / 300 (59.1%) 예요. 약 30% recall에서 60% 미만으로 급격히 개선됐어요. 아주 단순한 접근이었지만 DSPy는 여기서 계속 반복할 많은 도구를 제공해요.

최적화된 프롬프트를 살펴보며 뭘 배웠는지 이해해 볼게요. 쿼리 하나를 실행한 뒤 마지막 두 프롬프트를 확인하면 Hop() 프로그램의 후반 반복에서 두 하위 모듈에 쓰인 프롬프트가 보여요. (MLflow 추적을 켰다면 에이전트가 한 모든 단계를 풍부한 트리 뷰로 볼 수 있어요.)

optimized(claim="The author of the 1960s unproduced script written for The Beatles, Up Against It, and Bernard-Marie Koltès are both playwrights.").titles

결과는 ['Up Against It', 'Bernard-Marie Koltès', 'The Beatles', 'Joe Orton']예요.

dspy.inspect_history(n=2)로 두 하위 모듈의 프롬프트를 확인할 수 있어요.

generate_query 서브모듈은 claim과 notes를 입력받아 query를 만들어요. 최적화된 지시문은 "주장과 메모 집합이 주어지면, 주장을 뒷받침하거나 반박할 추가 증거·문맥을 모을 수 있는 쿼리를 만들라. 단계적으로 생각해서 메모의 정보와 관련된 구체적·관련성 높은 쿼리를 만들라"예요. 예를 들어 Chen Xiuke 출생지 문제에서는 메모에서 출생지가 Dongfang, Hainan임을 확인하고 "What is the birthplace of Chen Xiuke?" 쿼리를 만들어요. notes가 비어 있으면 Beatle 스크립트 저자를 묻는 쿼리를 만들어요.

append_notes 서브모듈은 claim, notes, context를 입력받아 new_notes와 titles를 만들어요. 최적화된 지시문은 "주장·메모·문맥을 분석해 주장을 뒷받침하거나 반박할 새 메모를 만들고, 문맥에서 핵심 주제·엔티티를 나타내는 관련 제목을 추출하라"예요. 예를 들어 Zak Ové 아버지에 대한 주장에서는 문맥에서 Horace Ové가 사진작가·영화감독·작가임을 확인하고, A. Edward Sutherland는 사진작가가 아니라 영화감독이므로 주장이 맞다고 결론을 내려요. new_notes에 해당 사실들을 기록하고 titles로 Horace Ové, A. Edward Sutherland, Zak Ové를 내놓아요.

마지막으로 최적화된 프로그램을 저장해 나중에 다시 쓸게요.

optimized.save("optimized_hop.json")

loaded_program = Hop()
loaded_program.load("optimized_hop.json")

loaded_program(claim="The author of the 1960s unproduced script written for The Beatles, Up Against It, and Bernard-Marie Koltès are both playwrights.").titles

로드된 프로그램도 같은 질문에 ['Up Against It', 'Bernard-Marie Koltès', 'The Beatles', 'Joe Orton']을 돌려줘요.

MLflow에 프로그램 저장

로컬 파일 대신 MLflow에 저장하면 재현성과 협업에 유리해요. 1) MLflow가 고정된 환경 메타데이터를 프로그램과 함께 자동 저장하고, 2) 프로그램의 성능·비용을 함께 추적하며, 3) 실험 공유로 팀과 결과를 나눌 수 있어요.

import mlflow

# Start an MLflow Run and save the program
with mlflow.start_run(run_name="optimized"):
    model_info = mlflow.dspy.log_model(
        optimized,
        artifact_path="model", # Any name to save the program in MLflow
    )

# Load the program back from MLflow
loaded = mlflow.dspy.load_model(model_info.model_uri)

더 알아보기 (Learn more)