엔티티 추출 튜토리얼

엔티티 추출 튜토리얼 (Entity Extraction)

이 튜토리얼은 CoNLL-2003 데이터셋으로 **엔티티 추출(entity extraction)**을 수행하는 방법을 보여줘요. 특히 '사람'을 가리키는 엔티티를 추출하는 데 초점을 맞춰요.

  • CoNLL-2003 데이터셋에서 사람을 가리키는 엔티티를 추출·라벨링
  • 사람 엔티티를 추출하는 DSPy 프로그램 정의
  • CoNLL-2003 데이터셋 일부에서 프로그램을 최적화하고 평가

튜토리얼을 마치면 DSPy에서 signature와 모듈로 작업을 구조화하는 법, 시스템 성능을 평가하는 법, 옵티마이저로 품질을 높이는 법을 이해하게 돼요. 최신 DSPy를 설치하고 따라오시고, DSPy에 대한 개념적 개요가 더 필요하다면 최근 강의를 추천해요.

# Install the latest version of DSPy
%pip install -U dspy
# Install the Hugging Face datasets library to load the CoNLL-2003 dataset
%pip install datasets

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

출처: Entity Extraction

MLflow DSPy 통합

MLflow는 DSPy와 기본적으로 통합되는 LLMOps 도구로, 설명 가능성과 실험 추적을 제공해요. 네 단계로 쉽게 설정할 수 있어요.

  1. MLflow 설치
%pip install mlflow>=2.20
  1. 별도 터미널에서 MLflow UI 시작
mlflow ui --port 5000
  1. 노트북을 MLflow에 연결
import mlflow

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("DSPy")
  1. 추적 활성화
mlflow.dspy.autolog()

자세한 내용은 MLflow DSPy 문서를 참고하세요.

데이터셋 로드와 준비

CoNLL-2003은 엔티티 추출 작업에 흔히 쓰이는 데이터셋이에요. 토큰들에 사람·조직·위치 같은 엔티티 라벨이 붙어 있어요. 우리는 1) Hugging Face datasets 라이브러리로 데이터셋을 로드하고, 2) 사람 토큰을 뽑는 함수를 정의하고, 3) 훈련·테스트용 작은 부분집합으로 슬라이스할 거예요. DSPy는 구조화된 형식의 예시를 기대하니, 데이터셋을 DSPy Example로 변환해 줄게요.

import os
import tempfile
from datasets import load_dataset
from typing import Dict, Any, List
import dspy

def load_conll_dataset() -> dict:
    """
    Loads the CoNLL-2003 dataset into train, validation, and test splits.
    
    Returns:
        dict: Dataset splits with keys 'train', 'validation', and 'test'.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        # Use a temporary Hugging Face cache directory for compatibility with certain hosted notebook
        # environments that don't support the default Hugging Face cache directory
        os.environ["HF_DATASETS_CACHE"] = temp_dir
        return load_dataset("lhoestq/conll2003")

def extract_people_entities(data_row: dict[str, Any]) -> list[str]:
    """
    Extracts entities referring to people from a row of the CoNLL-2003 dataset.
    
    Args:
        data_row (dict[str, Any]): A row from the dataset containing tokens and NER tags.
    
    Returns:
        list[str]: List of tokens tagged as people.
    """
    return [
        token
        for token, ner_tag in zip(data_row["tokens"], data_row["ner_tags"])
        if ner_tag in (1, 2)  # CoNLL entity codes 1 and 2 refer to people
    ]

def prepare_dataset(data_split, start: int, end: int) -> list[dspy.Example]:
    """
    Prepares a sliced dataset split for use with DSPy.
    
    Args:
        data_split: The dataset split (e.g., train or test).
        start (int): Starting index of the slice.
        end (int): Ending index of the slice.
    
    Returns:
        list[dspy.Example]: List of DSPy Examples with tokens and expected labels.
    """
    return [
        dspy.Example(
            tokens=row["tokens"],
            expected_extracted_people=extract_people_entities(row)
        ).with_inputs("tokens")
        for row in data_split.select(range(start, end))
    ]

# Load the dataset
dataset = load_conll_dataset()

# Prepare the training and test sets
train_set = prepare_dataset(dataset["train"], 0, 50)
test_set = prepare_dataset(dataset["test"], 0, 200)

extract_people_entities는 각 토큰의 ner_tag가 CoNLL 코드 1이나 2(사람을 가리킴)일 때 그 토큰을 뽑아요. 훈련 50개, 테스트 200개로 준비했어요.

DSPy 설정과 엔티티 추출 프로그램 만들기

사람 엔티티를 추출하는 DSPy 프로그램을 정의할게요.

여기서 소개되는 핵심 개념:

  • Signatures: 프로그램의 구조화된 입력/출력 스키마를 정의
  • Modules: 프로그램 로직을 재사용·조합 가능한 단위로 캡슐화

구체적으로는 PeopleExtraction signature로 입력(tokens)과 출력(extracted_people) 필드를 지정하고, dspy.ChainOfThought 모듈로 이 signature를 구현하는 people_extractor 프로그램을 만들어요. 그리고 dspy.LM 클래스와 dspy.configure()로 프로그램이 쓸 언어 모델을 설정해요.

from typing import List

class PeopleExtraction(dspy.Signature):
    """
    Extract contiguous tokens referring to specific people, if any, from a list of string tokens.
    Output a list of tokens. In other words, do not combine multiple tokens into a single value.
    """
    tokens: list[str] = dspy.InputField(desc="tokenized text")
    extracted_people: list[str] = dspy.OutputField(desc="all tokens referring to specific people extracted from the tokenized text")

people_extractor = dspy.ChainOfThought(PeopleExtraction)

DSPy가 OpenAI의 gpt-4o-mini를 쓰도록 설정할게요. 인증은 OPENAI_API_KEY에서 읽어요. 다른 프로바이더나 로컬 모델로도 쉽게 바꿀 수 있어요.

lm = dspy.LM(model="openai/gpt-4o-mini")
dspy.configure(lm=lm)

metric과 평가 함수 정의

DSPy에서 프로그램 성능을 평가하는 건 반복 개발의 핵심이에요. 좋은 평가 체계는 출력 품질을 측정하고, 정답 라벨과 비교하며, 개선할 부분을 찾게 해줘요.

  • 커스텀 metric(extraction_correctness_metric)으로 추출 엔티티가 정답과 일치하는지 평가
  • 평가 함수(evaluate_correctness)로 훈련·테스트 데이터셋에 metric을 적용해 전체 정확도 계산

평가 함수는 DSPy의 Evaluate 유틸리티로 병렬 처리와 결과 시각화를 처리해요.

def extraction_correctness_metric(example: dspy.Example, prediction: dspy.Prediction, trace=None) -> bool:
    """
    Computes correctness of entity extraction predictions.
    
    Args:
        example (dspy.Example): The dataset example containing expected people entities.
        prediction (dspy.Prediction): The prediction from the DSPy people extraction program.
        trace: Optional trace object for debugging.
    
    Returns:
        bool: True if predictions match expectations, False otherwise.
    """
    return prediction.extracted_people == example.expected_extracted_people

evaluate_correctness = dspy.Evaluate(
    devset=test_set,
    metric=extraction_correctness_metric,
    num_threads=24,
    display_progress=True,
    display_table=True
)

초기 추출기 평가

최적화 전에 기준선(baseline) 평가를 해야 해요. 그러면 최적화 후 비교 기준점이 생기고, 초기 구현의 약점도 드러나요. people_extractor 프로그램을 테스트 세트에서 실행하고 앞서 정의한 평가 체계로 정확도를 측정할게요.

evaluate_correctness(people_extractor, devset=test_set)

결과는 Average Metric: 172.00 / 200 (86.0%) 예요. 평가 표를 보면 첫 행에서 "JAPAN"과 "CHINA"를 사람으로 잘못 추출해(정답은 CHINA뿐) 오답 표시가 없는 걸 볼 수 있어요. 사람만 추출해야 하는데 국가 이름을 사람으로 오인하는 실수가 있죠.

MLflow에 평가 결과 기록하기

import mlflow

with mlflow.start_run(run_name="extractor_evaluation"):
    evaluate_correctness = dspy.Evaluate(
        devset=test_set,
        metric=extraction_correctness_metric,
        num_threads=24,
        display_progress=True,
    )

    # Evaluate the program as usual
    result = evaluate_correctness(people_extractor)

    # Log the aggregated score
    mlflow.log_metric("exact_match", result.score)
    # Log the detailed evaluation results as a table
    mlflow.log_table(
        {
            "Tokens": [example.tokens for example in test_set],
            "Expected": [example.expected_extracted_people for example in test_set],
            "Predicted": [output[1] for output in result.results],
            "Exact match": [output[2] for output in result.results],
        },
        artifact_file="eval_results.json",
    )

모델 최적화

DSPy는 시스템 품질을 높이는 강력한 옵티마이저를 제공해요. 여기서는 MIPROv2 옵티마이저를 쓸게요. 이 옵티마이저는 1) LM으로 프롬프트의 지시문을 조정하고, 2) dspy.ChainOfThought로 생성한 reasoning을 더해 훈련 데이터셋에서 few-shot 예시를 만들어 프롬프트를 자동으로 튜닝해요. 훈련 세트에서 정확도를 최대화하도록 최적화돼요.

mipro_optimizer = dspy.MIPROv2(
    metric=extraction_correctness_metric,
    auto="medium",
)
optimized_people_extractor = mipro_optimizer.compile(
    people_extractor,
    trainset=train_set,
    max_bootstrapped_demos=4,
    minibatch=False
)

최적화된 프로그램 평가

최적화 후 테스트 세트에서 다시 평가해 개선을 확인할게요. 초기 결과와 비교하면 최적화의 이점을 정량화하고, 프로그램이 보지 못한 데이터로 잘 일반화되는지 검증할 수 있어요. 이 경우 테스트 데이터셋에서 정확도가 크게 올랐어요.

evaluate_correctness(optimized_people_extractor, devset=test_set)

결과는 Average Metric: 186.00 / 200 (93.0%) 예요. 앞서 국가 이름을 사람으로 오인하던 오류가 대부분 사라졌어요.

최적화된 프로그램의 프롬프트 확인

inspect_history(n=1)로 마지막 상호작용을 보고 few-shot 예시가 어떻게 추가됐는지 확인할게요.

dspy.inspect_history(n=1)

최적화된 프롬프트의 system message를 보면 지시문이 더 구체적으로 다듬어져 있어요. "규제 준수와 공중 보건 커뮤니케이션에 정확한 개인 식별이 중요한 고위험 상황"이라는 문맥을 붙이고, 각 개인을 별개 토큰으로 출력하라고 강조해요. 그 아래에는 few-shot 예시들이 실려 있어요.

  • "European Union"만 언급된 토큰 → 조직이라 사람 없음으로 []
  • "BRUSSELS"와 날짜만 → 위치·날짜뿐이라 []
  • "Fischler"(사람) → ["Fischler"]
  • "Werner Zwingmann"(독일 대표) → ["Werner", "Zwingmann"]

이렇게 국가·조직을 사람으로 오인하지 않는 예시가 더해져 정확도가 개선된 거예요.

비용 관리하기

DSPy는 프로그램의 비용을 추적할 수 있어요. 지금까지 DSPy 추출기 프로그램이 만든 모든 LM 호출의 비용(USD)은 이렇게 구할 수 있어요.

cost = sum([x['cost'] for x in lm.history if x['cost'] is not None])  # cost in USD, as calculated by LiteLLM for certain providers
cost

예시 기준 약 0.26 USD 정도가 나와요.

최적화된 프로그램 저장과 로딩

DSPy는 프로그램 저장·로딩을 지원해서, 최적화된 시스템을 처음부터 다시 최적화하지 않고 재사용할 수 있어요. 프로덕션 배포나 동료와 공유할 때 특히 유용하죠.

optimized_people_extractor.save("optimized_extractor.json")

loaded_people_extractor = dspy.ChainOfThought(PeopleExtraction)
loaded_people_extractor.load("optimized_extractor.json")

loaded_people_extractor(tokens=["Italy", "recalled", "Marcello", "Cuttitta"]).extracted_people

결과는 ['Marcello', 'Cuttitta']예요. 저장된 프로그램을 새로 로드해도 잘 동작해요.

MLflow에 프로그램 저장

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

import mlflow

# Start an MLflow Run and save the program
with mlflow.start_run(run_name="optimized_extractor"):
    model_info = mlflow.dspy.log_model(
        optimized_people_extractor,
        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)

결론

이 튜토리얼에서는 DSPy로 모듈화되고 해석 가능한 엔티티 추출 시스템을 만들고, DSPy 내장 도구로 평가·최적화하는 방법을 살펴봤어요. 구조화된 입출력을 활용해 시스템을 이해하고 개선하기 쉽게 만들었고, 최적화 과정에서 프롬프트를 손으로 짜거나 파라미터를 만지작거리지 않고도 성능을 빠르게 끌어올릴 수 있었어요.

다음 단계:

  • 다른 엔티티 유형(위치, 조직 등) 추출 실험
  • 더 복잡한 추론 작업에 ReAct 같은 DSPy의 다른 내장 모듈 탐색
  • 대규모 문서 처리·요약 같은 더 큰 워크플로에 시스템 적용

더 알아보기 (Learn more)