Finetuning API와 LlamaIndex로 MistralAI 모델 파인튜닝하기

Finetuning API와 LlamaIndex로 MistralAI 모델 파인튜닝하기

MistralAI 파인튜닝 API로 open-mistral-7b를 파인튜닝하는 예제노트북이에요. mistral-large-latest로 훈련 데이터를 생성해 지식을 증류(distill)하고, 그 데이터로 open-mistral-7b를 파인튜닝합니다. 평가는 ragas 라이브러리로 수행하고, 메트릭은 Weights & Biases에서 모니터링할 수 있어요.

출처: 문서

본문

구체적으로 mistral-large-latest로 훈련 데이터를 생성해 open-mistral-7b를 파인튜닝함으로써 mistral-large-latest의 지식을 증류하려고 합니다. 모든 훈련 데이터는 인덱스 데이터의 서로 다른 두 섹션으로 생성되어 훈련 세트와 평가 세트를 만듭니다. 합성 훈련·평가 질문은 mistral-small-latest를 사용해 open-mistral-7b와 mistral-large-latest에 대한 편향을 피합니다. 그런 다음 MistralAIFinetuneEngine 래퍼 추상화로 파인튜닝합니다.

필요한 패키지를 설치하고 설정합니다.

%pip install llama-index-finetuning
%pip install llama-index-finetuning-callbacks
%pip install llama-index-llms-mistralai
%pip install llama-index-embeddings-mistralai
# !pip install llama-index pypdf sentence-transformers ragas
# NESTED ASYNCIO LOOP NEEDED TO RUN ASYNC IN A NOTEBOOK
import nest_asyncio
nest_asyncio.apply()
import os
os.environ["MISTRAL_API_KEY"] = "<YOUR MISTRALAI API KEY>"

데이터 다운로드

먼저 훈련 데이터를 생성할 PDF를 내려받습니다.

!curl https://www.ipcc.ch/report/ar6/wg2/downloads/report/IPCC_AR6_WGII_Chapter03.pdf --output IPCC_AR6_WGII_Chapter03.pdf

훈련·평가 데이터 생성

다음 단계는 훈련 및 평가 데이터셋을 생성하는 것입니다. 내려받은 PDF의 서로 다른 섹션에 대해 훈련 질문 40개와 평가 질문 40개를 만들게요. 평가 질문에는 open-mistral-7b를 사용해 기준(baseline) 성능을 얻습니다. 그런 다음 훈련 질문에는 mistral-large-latest를 사용해 훈련 데이터를 생성합니다.

from llama_index.core import SimpleDirectoryReader
from llama_index.core.evaluation import DatasetGenerator

documents = SimpleDirectoryReader(
    input_files=["IPCC_AR6_WGII_Chapter03.pdf"]
).load_data()
from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding

open_mistral = MistralAI(
    model="open-mistral-7b", temperature=0.1
)  # model to be finetuning

mistral_small = MistralAI(
    model="mistral-small-latest", temperature=0.1
)  # model for question generation

embed_model = MistralAIEmbedding()
question_gen_query = (
    "You are a Teacher/ Professor. Your task is to setup "
    "a quiz/examination. Using the provided context, formulate "
    "a single question that captures an important fact from the "
    "context. Restrict the question to the context information provided."
    "You should generate only question and nothing else."
)

dataset_generator = DatasetGenerator.from_documents(
    documents[:80],
    question_gen_query=question_gen_query,
    llm=mistral_small,
)

훈련 질문 40개와 평가 질문 40개를 생성합니다.

# Note: This might take sometime.
questions = dataset_generator.generate_questions_from_nodes(num=40)
print("Generated ", len(questions), " questions")
questions[10:15]
with open("train_questions.txt", "w") as f:
    for question in questions:
        f.write(question + "\n")

이제 평가 데이터셋을 만들기 위해 완전히 다른 문서 세트에서 질문을 생성합니다.

dataset_generator = DatasetGenerator.from_documents(
    documents[80:],
    question_gen_query=question_gen_query,
    llm=mistral_small,
)
with open("eval_questions.txt", "w") as f:
    for question in questions:
        f.write(question + "\n")

open-mistral-7b 쿼리 엔진으로 초기 평가

이 평가에는 ragas 평가 라이브러리를 사용합니다. Ragas는 RAG 파이프라인을 위한 많은 평가 메트릭을 제공합니다. 이 노트북에서는 다음 두 메트릭을 사용할게요.

  • answer_relevancy: 생성된 답변이 프롬프트에 얼마나 관련 있는지 측정합니다. 답변이 불완전하거나 중복 정보를 담으면 점수가 낮아집니다. LLM이 생성된 답변으로 주어진 질문을 생성할 확률로 정량화하며, 값은 (0, 1) 범위로 높을수록 좋습니다.
  • faithfulness: 생성된 답변의 주어진 문맥 대비 사실적 일관성을 측정합니다. 생성된 답변에서 문장을 만들고 각 문장을 문맥에 대해 검증하는 다단계 패러다임으로 수행됩니다. 답변은 (0, 1) 범위로 스케일되며 높을수록 좋습니다.
questions = []
with open("eval_questions.txt", "r") as f:
    for line in f:
        questions.append(line.strip())
from llama_index.core import VectorStoreIndex
from llama_index.embeddings.mistralai import MistralAIEmbedding

# limit the context window to 2048 tokens so that refine is used
from llama_index.core import Settings
Settings.context_window = 2048
Settings.llm = open_mistral
Settings.embed_model = MistralAIEmbedding()

index = VectorStoreIndex.from_documents(
    documents,
)
query_engine = index.as_query_engine(similarity_top_k=2)
contexts = []
answers = []

for question in questions:
    response = query_engine.query(question)
    contexts.append([x.node.get_content() for x in response.source_nodes])
    answers.append(str(response))

평가는 RAGAS를 위해 OpenAI LLM을 사용합니다.

# We will use OpenAI LLM for evaluation using RAGAS
os.environ["OPENAI_API_KEY"] = "<YOUR OPENAI API KEY>"
from datasets import Dataset
from ragas import evaluate
from ragas.metrics import answer_relevancy, faithfulness

ds = Dataset.from_dict(
    {
        "question": questions,
        "answer": answers,
        "contexts": contexts,
    }
)
result = evaluate(ds, [answer_relevancy, faithfulness])

파인튜닝 전의 결과를 확인합니다.

print(result)

mistral-large-latest로 훈련 데이터 수집

여기서는 open-mistral-7b가 파인튜닝할 데이터를 수집하기 위해 mistral-large-latest를 사용합니다.

from llama_index.llms.mistralai import MistralAI
from llama_index.finetuning.callbacks import MistralAIFineTuningHandler
from llama_index.core.callbacks import CallbackManager

finetuning_handler = MistralAIFineTuningHandler()
callback_manager = CallbackManager([finetuning_handler])

llm = MistralAI(model="mistral-large-latest", temperature=0.1)
llm.callback_manager = callback_manager
questions = []
with open("train_questions.txt", "r") as f:
    for line in f:
        questions.append(line.strip())
from llama_index.core import VectorStoreIndex

Settings.embed_model = MistralAIEmbedding()
Settings.llm = llm

index = VectorStoreIndex.from_documents(
    documents,
)
query_engine = index.as_query_engine(similarity_top_k=2)
from tqdm import tqdm

contexts = []
answers = []

for question in tqdm(questions, desc="Processing questions"):
    response = query_engine.query(question)
    contexts.append(
        "\n".join([x.node.get_content() for x in response.source_nodes])
    )
    answers.append(str(response))

질문·문맥·답변을 Mistral 파인튜닝 API가 요구하는 JSONL 형식으로 변환합니다.

import json
from typing import List

def convert_data_jsonl_format(
    questions: List[str],
    contexts: List[str],
    answers: List[str],
    output_file: str,
) -> None:
    with open(output_file, "w") as outfile:
        for context, question, answer in zip(contexts, questions, answers):
            message_dict = {
                "messages": [
                    {
                        "role": "assistant",
                        "content": "You are a helpful assistant to answer user queries based on provided context.",
                    },
                    {
                        "role": "user",
                        "content": f"context: {context} \n\n question: {question}",
                    },
                    {"role": "assistant", "content": answer},
                ]
            }
            # Write the JSON object in a single line
            json.dump(message_dict, outfile)
            # Add a newline character after each JSON object
            outfile.write("\n")

convert_data_jsonl_format(questions, contexts, answers, "training.jsonl")

MistralAIFinetuneEngine 만들기

MistralAIFinetuneEngine을 생성합니다. 이 파인튜닝 엔진은 파인튜닝 작업을 실행하고, 나머지 LlamaIndex 워크플로에 바로 연결할 수 있는 LLM 모델을 반환해 줍니다. 여기서는 기본 생성자를 쓰지만, from_finetuning_handler 클래스 메서드로 파인튜닝 핸들러를 직접 전달할 수도 있어요.

from llama_index.llms.mistralai import MistralAI
from llama_index.finetuning.callbacks import MistralAIFineTuningHandler
from llama_index.core.callbacks import CallbackManager
from llama_index.finetuning.mistralai import MistralAIFinetuneEngine

# Wandb for monitorning the training logs
wandb_integration_dict = {
    "project": "mistralai",
    "run_name": "finetuning",
    "api_key": "3a486c0d4066396b07b0ebf6826026c24e981f37",
}

finetuning_engine = MistralAIFinetuneEngine(
    base_model="open-mistral-7b",
    training_path="training.jsonl",
    # validation_path="<validation file>", # validation file is optional
    verbose=True,
    training_steps=5,
    learning_rate=0.0001,
    wandb_integration_dict=wandb_integration_dict,
)

파인튜닝을 실행합니다.

# starts the finetuning of open-mistral-7b
finetuning_engine.finetune()

작업의 현재 상태('RUNNING')를 확인합니다.

# This will show the current status of the job - 'RUNNING'
finetuning_engine.get_current_job()

작업의 현재 상태('SUCCESS')를 확인합니다.

# This will show the current status of the job - 'SUCCESS'
finetuning_engine.get_current_job()
ft_llm = finetuning_engine.get_finetuned_model(temperature=0.1)

평가 (Evaluation)

파인튜닝된 모델이 생성되면 평가 데이터셋으로 다시 실행해 성능 향상을 측정합니다.

from llama_index.core import Settings
from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding

# Setting up finetuned llm
Settings.llm = ft_llm
Settings.context_window = (
    2048  # limit the context window artifically to test refine process
)
Settings.embed_model = MistralAIEmbedding()
from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine(similarity_top_k=2, llm=ft_llm)
contexts = []
answers = []

for question in tqdm(questions, desc="Processing Questions"):
    response = query_engine.query(question)
    contexts.append([x.node.get_content() for x in response.source_nodes])
    answers.append(str(response))
ds = Dataset.from_dict(
    {
        "question": questions,
        "answer": answers,
        "contexts": contexts,
    }
)
result = evaluate(ds, [answer_relevancy, faithfulness])

파인튜닝된 모델의 결과를 확인합니다.

open-mistral-7b : 'answer_relevancy': 0.8151, 'faithfulness': 0.8360
open-mistral-7b-finetuned : 'answer_relevancy': 0.8016, 'faithfulness': 0.8924

보시다시피 faithfulness 점수는 증가했고 answer relevancy는 약간 떨어졌습니다. 파인튜닝이 실제로 뭔가를 바꿨는지 확인하기 위해 응답의 차이를 빠르게 비교해 봅니다.

차이 탐구 (Exploring Differences)

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)
print(questions[20])
from llama_index.core.response.notebook_utils import display_response
query_engine = index.as_query_engine(llm=open_mistral)
response = query_engine.query(questions[20])
display_response(response)
query_engine = index.as_query_engine(llm=ft_llm)
response = query_engine.query(questions[20])
display_response(response)

보시다시피 파인튜닝된 모델이 더 철저한 응답을 제공합니다! 이는 검색된 문맥을 답변이 더 잘 대표하므로 ragas의 향상된 faithfulness 점수와 일치해요.

결론 (Conclusion)

요약하자면, 약 40개의 질문만으로 파인튜닝한 것이 평가 점수를 개선하는 데 실제로 도움이 됐습니다.

  • answer_relevancy: 0.8151 → 0.8016 (약간 하락하지만 매우 작음)
  • faithfulness: 0.8360 → 0.8924 (개선됨)

faithfulness가 개선된 것은 주어진 원래 질문에 답변이 더 잘 부합한다는 뜻입니다.

더 알아보기 (Learn more)

  • Mistral Fine-tuning 문서 — 모델 파인튜닝 가이드
  • Ragas 문서 — RAG 평가 라이브러리
  • MistralAIFinetuneEngine / MistralAIFineTuningHandler — LlamaIndex 파인튜닝 구성 요소
  • 모델: open-mistral-7b(파인튜닝 대상), mistral-small-latest(질문 생성), mistral-large-latest(훈련 데이터 생성)