튜토리얼: 엔티티 추출

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

이 튜토리얼에서는 DSPy로 CoNLL-2003 데이터셋을 사용해 **엔티티 추출(entity extraction)**을 수행하는 방법을 보여드릴게요. 초점은 사람을 가리키는 엔티티를 추출하는 것입니다. 우리가 할 일은:

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

이 튜토리얼이 끝나면, DSPy에서 시그니처(signatures)와 모듈(modules)을 사용해 작업을 구조화하고, 시스템 성능을 평가하며, 옵티마이저로 품질을 개선하는 방법을 이해하게 될 거예요.

최신 버전의 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

%pip install mlflow>=2.20

mlflow ui --port 5000

import mlflow

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("DSPy")

mlflow.dspy.autolog()

데이터셋 로드 및 준비 (Load and Prepare the Dataset)

이 섹션에서는 엔티티 추출 작업에 흔히 사용되는 CoNLL-2003 데이터셋을 준비할게요. 데이터셋에는 사람, 조직, 위치 같은 엔티티 라벨로 주석이 달린 토큰이 포함됩니다.

우리가 할 일:

1. Hugging Face datasets 라이브러리로 데이터셋 로드하기.
2. 사람을 가리키는 토큰을 추출하는 함수 정의하기.
3. 훈련과 테스트를 위한 더 작은 부분집합을 만들도록 데이터셋 슬라이싱하기.

DSPy는 구조화된 형식의 예를 기대하므로, 손쉬운 통합을 위해 데이터셋을 DSPy Examples로 변환할 거예요.

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)

DSPy 구성 및 엔티티 추출 프로그램 만들기 (Configure DSPy and create an Entity Extraction Program)

여기서 토큰화된 텍스트에서 사람을 가리키는 엔티티를 추출하는 DSPy 프로그램을 정의할게요.

그런 다음 프로그램의 모든 호출에 특정 언어 모델(gpt-4o-mini)을 사용하도록 DSPy를 구성합니다.

소개되는 핵심 DSPy 개념:

• 시그니처(Signatures): 프로그램을 위한 구조화된 입력/출력 스키마를 정의해요.
• 모듈(Modules): 프로그램 로직을 재사용 가능하고 조합 가능한 단위로 캡슐화해요.

구체적으로 우리가 할 일:

• 입력(tokens)과 출력(extracted_people) 필드를 지정하는 PeopleExtraction DSPy 시그니처 만들기.
• DSPy의 내장 dspy.ChainOfThought 모듈을 사용해 PeopleExtraction 시그니처를 구현하는 people_extractor 프로그램 정의하기. 이 프로그램은 언어 모델(LM) 프롬프트를 사용해 입력 토큰 목록에서 사람을 가리키는 엔티티를 추출해요.
• dspy.LM 클래스와 dspy.configure() 메서드를 사용해 프로그램이 호출될 때 DSPy가 사용할 언어 모델을 구성하기.

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)

여기서 프로그램에 OpenAI의 gpt-4o-mini 모델을 사용하도록 DSPy에 지시해요. 인증을 위해 DSPy는 OPENAI_API_KEY를 읽습니다. 다른 제공자나 로컬 모델로 쉽게 바꿀 수 있어요.

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

지표 및 평가 함수 정의 (Define Metric and Evaluation Functions)

DSPy에서 프로그램 성능을 평가하는 것은 반복적 개발에 중요해요. 좋은 평가 프레임워크는 다음을 가능하게 합니다:

• 프로그램 출력의 품질 측정.
• 출력을 정답 라벨과 비교.
• 개선 영역 식별.

우리가 할 일:

• 추출된 엔티티가 정답과 일치하는지 평가하는 커스텀 지표(extraction_correctness_metric) 정의하기.
• 이 지표를 훈련/테스트 데이터셋에 적용하고 전체 정확도를 계산하는 평가 함수(evaluate_correctness) 만들기.

평가 함수는 결과의 병렬 처리와 시각화를 다루기 위해 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
)

초기 추출기 평가 (Evaluate Initial Extractor)

프로그램을 최적화하기 전에, 현재 성능을 이해하기 위한 기준(baseline) 평가가 필요해요. 이는 다음을 도와줍니다:

• 최적화 후 비교를 위한 기준점 설정.
• 초기 구현의 잠재적 약점 식별.

이 단계에서는 people_extractor 프로그램을 테스트 세트에서 실행하고 앞서 정의한 평가 프레임워크를 사용해 정확도를 측정할 거예요.

evaluate_correctness(people_extractor, devset=test_set)

Average Metric: 172.00 / 200 (86.0%): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:16<00:00, 11.94it/s]
2024/11/18 21:08:04 INFO dspy.evaluate.evaluate: Average Metric: 172 / 200 (86.0%)

86.0

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",
    )

모델 최적화 (Optimize the Model)

DSPy에는 시스템의 품질을 개선할 수 있는 강력한 옵티마이저가 포함되어 있어요.

여기서는 DSPy의 MIPROv2 옵티마이저를 사용해 다음을 수행합니다:

• 프로그램의 언어 모델(LM) 프롬프트를 자동으로 튜닝 — 1) LM을 사용해 프롬프트의 지침을 조정하고, 2) dspy.ChainOfThought에서 생성된 추론으로 증강된 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 Optimized Program)

최적화 후, 개선을 측정하기 위해 프로그램을 테스트 세트에서 다시 평가해요. 최적화된 결과와 초기 결과를 비교하면 다음이 가능합니다:

• 최적화의 이점을 정량화.
• 프로그램이 보이지 않는 데이터에 잘 일반화되는지 검증.

이 경우 테스트 데이터셋에서 프로그램의 정확도가 크게 개선된 것을 볼 수 있어요.

evaluate_correctness(optimized_people_extractor, devset=test_set)

Average Metric: 186.00 / 200 (93.0%): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:23<00:00,  8.58it/s]
2024/11/18 21:15:00 INFO dspy.evaluate.evaluate: Average Metric: 186 / 200 (93.0%)

93.0

최적화된 프로그램의 프롬프트 검사 (Inspect Optimized Program's Prompt)

프로그램을 최적화한 후, 상호작용 히스토리를 검사해 DSPy가 few-shot 예제로 프로그램의 프롬프트를 어떻게 증강했는지 확인할 수 있어요. 이 단계는 다음을 보여줍니다:

• 프로그램이 사용하는 프롬프트의 구조.
• 모델의 동작을 안내하기 위해 few-shot 예제가 어떻게 추가되는지.

inspect_history(n=1)을 사용해 마지막 상호작용을 보고 생성된 프롬프트를 분석해요.

dspy.inspect_history(n=1)

[2024-11-18T21:15:00.584497]

System message:

Your input fields are:
1. `tokens` (list[str]): tokenized text

Your output fields are:
1. `rationale` (str): ${produce the extracted_people}. We ...
2. `extracted_people` (list[str]): all tokens referring to specific people extracted from the tokenized text

...

In adhering to this structure, your objective is: 
        In a high-stakes situation where accurate identification of individuals is critical for regulatory compliance and public health communication, extract contiguous tokens referring to specific people from the provided list of string tokens. Ensure that you output each identified individual as separate tokens without combining multiple tokens into a single value. This task is essential for ensuring clarity and accountability in communications pertaining to EU regulations and health matters.

...

한 개인을 하나의 토큰으로 유지하는 등 few-shot 예제들이 어떻게 삽입되었는지, 그리고 최적화된 지침이 어떻게 상황을 구체화하는지 확인할 수 있어요.

비용 관리 (Keeping an eye on cost)

DSPy를 사용하면 프로그램의 비용을 추적할 수 있어요. 다음 코드는 지금까지 DSPy 추출기 프로그램이 만든 모든 LM 호출의 비용을 얻는 방법을 보여줘요.

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.26362742999999983

최적화된 프로그램 저장 및 로드 (Saving and Loading Optimized Programs)

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']

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)

결론 (Conclusion)

이 튜토리얼에서는 다음 방법을 시연했어요:

• DSPy를 사용해 엔티티 추출을 위한 모듈식이고 해석 가능한 시스템 구축하기.
• DSPy의 내장 도구를 사용해 시스템 평가 및 최적화하기.

구조화된 입력과 출력을 활용함으로써 시스템을 이해하고 개선하기 쉽게 만들었어요. 최적화 과정 덕분에 프롬프트를 수동으로 만들거나 매개변수를 조정하지 않고도 성능을 빠르게 개선할 수 있었습니다.

다음 단계:

• 다른 엔티티 타입(예: 위치나 조직) 추출 실험하기.
• 더 복잡한 추론 작업을 위해 ReAct 같은 DSPy의 다른 내장 모듈 탐색하기.
• 대규모 문서 처리나 요약 같은 더 큰 워크플로우에서 시스템 사용하기.

더 알아보기 (Learn more)

• Multihop Search 튜토리얼
• Optimize AI Program 튜토리얼
• 빌드 AI 프로그램 튜토리얼