W&B Models와 함께 Weave 사용하기

W&B Models와 함께 Weave를 사용하는 방법을 알아보세요

이 문서는 대화형 노트북입니다. 로컬에서 실행하거나 아래 링크를 사용할 수 있습니다:

• Google Colab에서 열기
• GitHub에서 소스 보기

​

사전 요구 사항

먼저 필요한 라이브러리를 설치하고, API 키를 설정하고, W&B에 로그인한 다음, 새 W&B 프로젝트를 생성하세요.

1. pip를 사용해 weave, pandas, unsloth, wandb, litellm, pydantic, torch, faiss-gpu를 설치합니다.

%%capture
!pip install weave wandb pandas pydantic litellm faiss-gpu
python
%%capture
!pip install unsloth
# 최신 nightly Unsloth도 설치하세요!
!pip uninstall unsloth -y && pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

2. 환경에 설정된 필요한 API 키를 추가합니다.

import os

from google.colab import userdata

os.environ["WANDB_API_KEY"] = userdata.get("WANDB_API_KEY")  # W&B Models and Weave
os.environ["OPENAI_API_KEY"] = userdata.get(
    "OPENAI_API_KEY"
)  # OpenAI - for retrieval embeddings
os.environ["GEMINI_API_KEY"] = userdata.get(
    "GEMINI_API_KEY"
)  # Gemini - 기본 채팅 모델용

3. W&B에 로그인한 다음 새 프로젝트를 만듭니다.

import pandas as pd
import wandb

import weave

wandb.login()

PROJECT = "weave-cookboook-demo"
ENTITY = "wandb-smle"

weave.init(ENTITY + "/" + PROJECT)

​

Models Registry에서 ChatModel을 다운로드하고 UnslothLoRAChatModel 구현하기

이 시나리오에서는 Llama-3.2 모델이 성능 최적화를 위해 Model Team에서 이미 unsloth 라이브러리로 파인튜닝되어 있으며, W&B Models Registry에서 사용할 수 있습니다. 이 step에서는 레지스트리에서 파인튜닝된 ChatModel을 조회해 RagModel과 호환되도록 weave.Model로 변환합니다.

아래에서 참조하는 RagModel은 완전한 RAG 애플리케이션으로 볼 수 있는 최상위 weave.Model입니다. 여기에는 ChatModel, 벡터 데이터베이스, 프롬프트가 포함됩니다. ChatModel 역시 weave.Model이며, W&B 레지스트리에서 아티팩트를 다운로드하는 코드를 포함합니다. ChatModel은 RagModel의 일부로서 다른 유형의 LLM 채팅 모델도 지원할 수 있도록 모듈식으로 변경할 수 있습니다. 자세한 내용은 Weave에서 모델 보기를 참고하세요.

ChatModel을 로드하려면 어댑터가 포함된 unsloth.FastLanguageModel 또는 peft.AutoPeftModelForCausalLM을 사용하면 앱에 효율적으로 인테그레이션할 수 있습니다. 레지스트리에서 모델을 다운로드한 후에는 model_post_init 방법을 사용해 초기화 및 예측 로직을 설정할 수 있습니다. 이 step에 필요한 코드는 레지스트리의 Use 탭에서 확인할 수 있으며, 구현에 직접 복사해 사용할 수 있습니다.

아래 코드는 W&B Models Registry에서 조회한 파인튜닝된 Llama-3.2 모델을 관리, 초기화, 사용하기 위한 UnslothLoRAChatModel 클래스를 정의합니다. UnslothLoRAChatModel은 최적화된 Inference를 위해 unsloth.FastLanguageModel을 사용합니다. model_post_init 방법은 모델 다운로드와 설정을 처리하고, predict 방법은 사용자 쿼리를 처리해 응답을 생성합니다. 사용 사례에 맞게 코드를 조정하려면 MODEL_REG_URL을 파인튜닝된 모델의 올바른 레지스트리 경로로 업데이트하고, 하드웨어나 requirements에 따라 max_seq_length 또는 dtype 같은 parameter를 조정하세요.”

from typing import Any

from pydantic import PrivateAttr
from unsloth import FastLanguageModel

import weave

class UnslothLoRAChatModel(weave.Model):
    """
    모델 이름 외에 더 많은 파라미터를 저장하고 버전 관리할 수 있도록 별도의 ChatModel 클래스를 정의합니다.
    특히 특정 파라미터로 인해 fine-tuning(로컬 또는 aaS)을 고려하는 경우에 유용합니다.
    """

    chat_model: str
    cm_temperature: float
    cm_max_new_tokens: int
    cm_quantize: bool
    inference_batch_size: int
    dtype: Any
    device: str
    _model: Any = PrivateAttr()
    _tokenizer: Any = PrivateAttr()

    def model_post_init(self, __context):
        # 레지스트리의 "Use" 탭에서 이 내용을 그대로 붙여넣을 수 있습니다
        run = wandb.init(project=PROJECT, job_type="model_download")
        artifact = run.use_artifact(f"{self.chat_model}")
        model_path = artifact.download()

        # unsloth 버전 (기본 2배 빠른 Inference 활성화)
        self._model, self._tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_path,
            max_seq_length=self.cm_max_new_tokens,
            dtype=self.dtype,
            load_in_4bit=self.cm_quantize,
        )
        FastLanguageModel.for_inference(self._model)

    @weave.op()
    async def predict(self, query: list[str]) -> dict:
        # add_generation_prompt = true - 생성 시 반드시 추가해야 합니다
        input_ids = self._tokenizer.apply_chat_template(
            query,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to("cuda")

        output_ids = self._model.generate(
            input_ids=input_ids,
            max_new_tokens=64,
            use_cache=True,
            temperature=1.5,
            min_p=0.1,
        )

        decoded_outputs = self._tokenizer.batch_decode(
            output_ids[0][input_ids.shape[1] :], skip_special_tokens=True
        )

        return "".join(decoded_outputs).strip()
python
MODEL_REG_URL = "wandb32/wandb-registry-RAG Chat Models/Finetuned Llama-3.2:v3"

max_seq_length = 2048  # 원하는 값을 선택하세요! 내부적으로 RoPE Scaling을 자동 지원합니다!
dtype = (
    None  # 자동 감지는 None. Tesla T4, V100은 Float16, Ampere 이상은 Bfloat16
)
load_in_4bit = True  # 메모리 사용량 절감을 위해 4bit 양자화 사용. False로 설정 가능.

new_chat_model = UnslothLoRAChatModel(
    name="UnslothLoRAChatModelRag",
    chat_model=MODEL_REG_URL,
    cm_temperature=1.0,
    cm_max_new_tokens=max_seq_length,
    cm_quantize=load_in_4bit,
    inference_batch_size=max_seq_length,
    dtype=dtype,
    device="auto",
)
python
await new_chat_model.predict(
    [{"role": "user", "content": "What is the capital of Germany?"}]
)

​

새 ChatModel 버전을 RagModel에 통합하기

파인튜닝된 채팅 모델로 RAG 애플리케이션을 구축하면 전체 파이프라인을 다시 빌드하지 않고도 맞춤형 컴포넌트를 활용해 대화형 AI를 개선할 수 있습니다. 이 step에서는 Weave 프로젝트에서 기존 RagModel을 조회하고, 새로 미세 조정한 모델을 사용하도록 해당 ChatModel을 업데이트합니다. 이렇게 원활하게 교체하면 벡터 데이터베이스(VDB)와 프롬프트 같은 다른 컴포넌트는 그대로 유지되므로, 애플리케이션의 전체 구조를 보존하면서 성능을 향상시킬 수 있습니다.

아래 코드는 Weave 프로젝트의 레퍼런스를 사용해 RagModel 객체를 조회합니다. 그런 다음 RagModel의 chat_model 속성을 업데이트해 이전 step에서 생성한 새 UnslothLoRAChatModel 인스턴스를 사용하도록 합니다. 그 후 업데이트된 RagModel을 게시하여 새 버전을 생성합니다. 마지막으로 업데이트된 RagModel을 사용해 샘플 예측 쿼리를 실행하여 새 채팅 모델이 사용되고 있는지 확인합니다.

RagModel = weave.ref(
    "weave://wandb-smle/weave-cookboook-demo/object/RagModel:cqRaGKcxutBWXyM0fCGTR1Yk2mISLsNari4wlGTwERo"
).get()
python
RagModel.chat_model.chat_model
python
await RagModel.predict("When was the first conference on climate change?")
python
# MAGIC: chat_model을 교체하고 새 버전을 게시합니다 (다른 RAG 컴포넌트는 신경 쓸 필요 없음)
RagModel.chat_model = new_chat_model
python
RagModel.chat_model.chat_model
python
# 예측 시 새 버전을 참조할 수 있도록 먼저 새 버전을 게시합니다
PUB_REFERENCE = weave.publish(RagModel, "RagModel")
python
await RagModel.predict("When was the first conference on climate change?")

​

weave.Evaluation 실행

다음 step에서는 기존 weave.Evaluation을 사용해 업데이트된 RagModel의 성능을 평가합니다. 이 과정은 새로 파인튜닝된 채팅 모델이 RAG 애플리케이션에서 예상대로 동작하는지 확인합니다. 인테그레이션을 간소화하고 Models 팀과 Apps 팀 간 협업을 지원하기 위해 평가 결과를 모델의 W&B run과 Weave workspace에 모두 로깅합니다.

In Models:

• 평가 summary는 파인튜닝된 채팅 모델을 다운로드하는 데 사용한 W&B run에 로깅됩니다. 여기에는 분석을 위해 workspace view에 표시되는 summary 메트릭과 그래프가 포함됩니다.
• 평가 트레이스 ID는 run의 설정에 추가되며, Model Team이 더 쉽게 추적할 수 있도록 Weave 페이지로 직접 연결됩니다.

In Weave:

• ChatModel의 아티팩트 또는 레지스트리 링크는 RagModel의 입력으로 저장됩니다.
• 더 풍부한 컨텍스트를 위해 W&B run ID가 평가 트레이스의 추가 column으로 저장됩니다.

아래 코드는 평가 객체를 조회하고, 업데이트된 RagModel을 사용해 평가를 실행한 뒤, 결과를 W&B와 Weave에 모두 로깅하는 방법을 보여줍니다. 평가 레퍼런스(WEAVE_EVAL)가 프로젝트 설정과 일치하는지 확인하세요.

# MAGIC: eval 데이터셋과 scorer를 사용하여 evaluation을 간단히 조회하고 활용할 수 있습니다
WEAVE_EVAL = "weave://wandb-smle/weave-cookboook-demo/object/climate_rag_eval:ntRX6qn3Tx6w3UEVZXdhIh1BWGh7uXcQpOQnIuvnSgo"
climate_rag_eval = weave.ref(WEAVE_EVAL).get()
python
with weave.attributes({"wandb-run-id": wandb.run.id}):
    # .call 속성을 사용하여 결과와 call을 모두 조회한 후 eval 트레이스를 Models에 저장합니다
    summary, call = await climate_rag_eval.evaluate.call(climate_rag_eval, RagModel)
python
# Models에 로그 기록
wandb.run.log(pd.json_normalize(summary, sep="/").to_dict(orient="records")[0])
wandb.run.config.update(
    {"weave_url": f"https://wandb.ai/wandb-smle/weave-cookboook-demo/r/call/{call.id}"}
)
wandb.run.finish()

​

새 RAG Model을 레지스트리에 저장

업데이트된 RagModel을 앞으로 Models 팀과 Apps 팀 모두가 사용할 수 있도록 W&B Models Registry에 레퍼런스 아티팩트로 푸시합니다.

아래 코드는 업데이트된 RagModel의 weave 객체 버전과 이름을 조회한 다음, 이를 사용해 레퍼런스 링크를 생성합니다. 이어서 모델의 Weave URL이 포함된 메타데이터와 함께 새 아티팩트를 W&B에 생성합니다. 이 아티팩트는 W&B Registry에 로깅되고 지정된 레지스트리 경로에 연결됩니다.

코드를 실행하기 전에 ENTITY 및 PROJECT 변수가 W&B 설정과 일치하는지, 그리고 대상 레지스트리 경로가 올바르게 지정되었는지 확인하세요. 이 과정은 새 RagModel을 W&B 생태계에 게시해 협업과 재사용을 쉽게 할 수 있도록 하여 워크플로를 마무리합니다.

MODELS_OBJECT_VERSION = PUB_REFERENCE.digest  # weave 객체 버전
MODELS_OBJECT_NAME = PUB_REFERENCE.name  # weave 객체 이름
python
models_url = f"https://wandb.ai/{ENTITY}/{PROJECT}/weave/objects/{MODELS_OBJECT_NAME}/versions/{MODELS_OBJECT_VERSION}"
models_link = (
    f"weave://{ENTITY}/{PROJECT}/object/{MODELS_OBJECT_NAME}:{MODELS_OBJECT_VERSION}"
)

with wandb.init(project=PROJECT, entity=ENTITY) as run:
    # 새 Artifact 생성
    artifact_model = wandb.Artifact(
        name="RagModel",
        type="model",
        description="Weave의 RagModel에서 가져온 Models 링크",
        metadata={"url": models_url},
    )
    artifact_model.add_reference(models_link, name="model", checksum=False)

    # 새 아티팩트 로그
    run.log_artifact(artifact_model, aliases=[MODELS_OBJECT_VERSION])

    # 레지스트리에 연결
    run.link_artifact(
        artifact_model, target_path="wandb32/wandb-registry-RAG Models/RAG Model"
    )