Weave 및 Models 통합 데모

Colab에서 실행해 보기 이 노트북은 W&B Weave를 W&B Models와 함께 사용하는 방법을 보여줍니다. 이 예제에서는 두 개의 서로 다른 팀을 가정합니다.

모델 팀: 모델 개발 팀이 새로운 Chat Model(Llama 3.2)을 파인튜닝하고 W&B Models를 사용해 레지스트리에 저장합니다.
앱 팀: 앱 개발 팀이 Chat Model을 가져와 W&B Weave를 사용해 새로운 RAG 챗봇을 생성하고 평가합니다.

W&B Models와 W&B Weave 둘 다에 대한 공개 워크스페이스는 여기에서 확인할 수 있습니다.

이 워크플로는 다음 단계를 포함합니다:

RAG 앱 코드를 W&B Weave로 계측합니다.
LLM(Llama 3.2와 같은 모델이지만, 원하는 다른 LLM으로 대체 가능)을 파인튜닝하고 W&B Models로 추적합니다.
파인튜닝된 모델을 W&B Registry에 로그합니다.
새로 파인튜닝한 모델로 RAG 앱을 구현하고, W&B Weave를 사용해 앱을 평가합니다.
결과에 만족하면, 업데이트된 RAG 앱에 대한 참조를 W&B Registry에 저장합니다.

참고: 아래에서 참조하는 RagModel은 최상위 weave.Model로, 하나의 완전한 RAG 앱으로 볼 수 있습니다. 이 모델은 ChatModel, 벡터 데이터베이스, 프롬프트를 포함합니다. ChatModel 역시 또 다른 weave.Model로, W&B Registry에서 아티팩트를 다운로드하는 코드를 포함하며, RagModel의 일부로서 다른 어떤 챗 모델도 지원할 수 있도록 변경할 수 있습니다. 자세한 내용은 Weave의 전체 모델을 참고하세요.

1. 설정

먼저 weave와 wandb를 설치한 후 API key로 로그인하세요. User Settings에서 API key를 생성하거나 조회할 수 있습니다.

pip install weave wandb

import wandb
import weave
import pandas as pd

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

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

2. 아티팩트를 기반으로 `ChatModel` 만들기

Registry에서 파인튜닝된 챗 모델을 가져와 weave.Model로 생성한 뒤, 다음 단계에서 RagModel에 바로 연결합니다. 이 모델은 기존 ChatModel과 동일한 파라미터를 사용하며, init과 predict 구현만 다릅니다.

pip install unsloth
pip uninstall unsloth -y && pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

모델 팀은 unsloth 라이브러리를 사용해 다양한 Llama-3.2 모델을 파인튜닝하여 더 빠르게 동작하도록 만들었습니다. 따라서 레지스트리에서 모델을 다운로드한 뒤 로드할 때에는 어댑터가 적용된 전용 unsloth.FastLanguageModel 또는 peft.AutoPeftModelForCausalLM 모델을 사용해야 합니다. 레지스트리의 “Use” 탭에서 모델 로딩 코드를 복사해 model_post_init에 붙여넣으세요.

import weave
from pydantic import PrivateAttr
from typing import Any, List, Dict, Optional
from unsloth import FastLanguageModel
import torch


class UnslothLoRAChatModel(weave.Model):
    """
    모델 이름 외에 더 많은 파라미터를 저장하고 버전 관리하기 위한 추가 ChatModel 클래스를 정의합니다.
    이를 통해 특정 파라미터에 대한 파인튜닝이 가능합니다.
    """

    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):
        run = wandb.init(project=PROJECT, job_type="model_download")
        artifact_ref = self.chat_model.replace("wandb-artifact://", "")
        artifact = run.use_artifact(artifact_ref)
        model_path = artifact.download()

        # unsloth 버전 (네이티브 2배 빠른 추론 활성화)
        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()

이제 레지스트리의 특정 링크를 사용해 새 모델을 생성하세요:

ORG_ENTITY = "wandb32"  # 여기에 조직 이름을 입력하세요
artifact_name = "Finetuned Llama-3.2" # 여기에 아티팩트 이름을 입력하세요
MODEL_REG_URL = f"wandb-artifact://{ORG_ENTITY}/wandb-registry-RAG Chat Models/{artifact_name}:v3"

max_seq_length = 2048
dtype = None
load_in_4bit = True

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

마지막으로 평가를 비동기적으로 실행하세요:

 await new_chat_model.predict(
     [{"role": "user", "content": "What is the capital of Germany?"}]
 )

3. 새로운 `ChatModel` 버전을 `RagModel`에 통합하기

파인튜닝된 챗 모델을 기반으로 RAG 앱을 구축하면, 특히 대화형 AI 시스템의 성능과 활용성을 크게 향상시킬 수 있는 여러 장점이 있습니다. 이제 기존 Weave 프로젝트에서 RagModel을 가져와서(RagModel의 현재 Weave ref는 아래 이미지와 같이 Use 탭에서 가져올 수 있습니다) ChatModel을 새로운 버전으로 교체합니다. 다른 컴포넌트(VDB, 프롬프트 등)는 변경하거나 다시 생성할 필요가 없습니다!

pip install litellm faiss-gpu

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

4. 기존 모델 실행에 연결되는 새로운 `weave.Evaluation` 실행

마지막으로 기존 weave.Evaluation에서 새로운 RagModel을 평가합니다. 통합을 최대한 쉽게 하기 위해 아래 변경 사항을 적용하세요. Models 관점에서:

레지스트리에서 모델을 가져오면 채팅 모델의 E2E 계보에 속하는 새로운 실행(run) 객체가 생성됩니다
모델 팀이 링크를 클릭해 해당 Weave 페이지로 이동할 수 있도록 Trace ID(현재 eval ID)를 실행 config에 추가합니다

Weave 관점에서:

아티팩트/레지스트리 링크를 ChatModel(즉, RagModel)의 입력으로 저장합니다
weave.attributes를 사용해 트레이스에 run.id를 추가 컬럼으로 저장합니다

# MAGIC: 평가 데이터셋과 스코어러가 포함된 평가를 가져와 사용합니다
WEAVE_EVAL = "weave://wandb-smle/weave-cookboook-demo/object/climate_rag_eval:ntRX6qn3Tx6w3UEVZXdhIh1BWGh7uXcQpOQnIuvnSgo"
climate_rag_eval = weave.ref(WEAVE_EVAL).get()

run = wandb.init()

with weave.attributes({"wandb-run-id": run.id}):
    # .call 속성을 사용하여 결과와 호출을 모두 검색하고 평가 트레이스를 Models에 저장합니다
    summary, call = await climate_rag_eval.evaluate.call(climate_rag_eval, ` RagModel `)

5. 새로운 RAG 모델을 Registry에 저장하기

새로운 RAG 모델을 효과적으로 공유하려면, 이를 참조용 아티팩트로 Registry에 푸시하고 Weave 버전을 별칭(alias)으로 추가합니다.

MODELS_OBJECT_VERSION = PUB_REFERENCE.digest  # weave 오브젝트 버전
MODELS_OBJECT_NAME = PUB_REFERENCE.name  # weave 오브젝트 이름

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_model = wandb.Artifact(
        name="RagModel",
        type="model",
        description="Models Link from RagModel in Weave",
        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"
    )

가이드

통합

튜토리얼

레퍼런스

1. 설정

2. 아티팩트를 기반으로 `ChatModel` 만들기

3. 새로운 `ChatModel` 버전을 `RagModel`에 통합하기

4. 기존 모델 실행에 연결되는 새로운 `weave.Evaluation` 실행

5. 새로운 RAG 모델을 Registry에 저장하기

가이드

통합

튜토리얼

레퍼런스

​1. 설정

​2. 아티팩트를 기반으로 ChatModel 만들기

​3. 새로운 ChatModel 버전을 RagModel에 통합하기

​4. 기존 모델 실행에 연결되는 새로운 weave.Evaluation 실행

​5. 새로운 RAG 모델을 Registry에 저장하기

1. 설정

2. 아티팩트를 기반으로 `ChatModel` 만들기

3. 새로운 `ChatModel` 버전을 `RagModel`에 통합하기

4. 기존 모델 실행에 연결되는 새로운 `weave.Evaluation` 실행

5. 새로운 RAG 모델을 Registry에 저장하기