Hugging Face Transformers

os.environ["WANDB_PROJECT"] = "<my-amazing-project>"  # W&B 프로젝트 이름 지정
os.environ["WANDB_LOG_MODEL"] = "checkpoint"  # 모든 모델 체크포인트 로깅

from transformers import TrainingArguments, Trainer

args = TrainingArguments(..., report_to="wandb")  # W&B 로깅 켜기
trainer = Trainer(..., args=args)

trainer.train()  # 트레이닝을 시작하고 W&B에 로깅

# 트레이닝 후 분석, 테스트, 기타 로깅된 코드

wandb.finish()

# 새 Run 만들기
with wandb.init(project="amazon_sentiment_analysis") as run:
    # Artifact 이름 및 버전 전달
    my_model_name = "model-bert-base-high-lr:latest"
    my_model_artifact = run.use_artifact(my_model_name)

    # 모델 가중치를 폴더에 다운로드하고 경로 반환
    model_dir = my_model_artifact.download()

    # 동일한 모델 클래스를 사용하여 해당 폴더에서 Hugging Face 모델 로드
    model = AutoModelForSequenceClassification.from_pretrained(
        model_dir, num_labels=num_labels
    )

    # 추가 트레이닝 또는 추론 실행

last_run_id = "xxxxxxxx"  # wandb Workspace에서 run_id 가져오기

# run_id에서 W&B Run 재개
with wandb.init(
    project=os.environ["WANDB_PROJECT"],
    id=last_run_id,
    resume="must",
) as run:
    # Artifact를 Run에 연결
    my_checkpoint_name = f"checkpoint-{last_run_id}:latest"
    my_checkpoint_artifact = run.use_artifact(my_model_name)

    # 체크포인트를 폴더에 다운로드하고 경로 반환
    checkpoint_dir = my_checkpoint_artifact.download()

    # 모델 및 Trainer 다시 초기화
    model = AutoModelForSequenceClassification.from_pretrained(
        "<model_name>", num_labels=num_labels
    )
    # 멋진 트레이닝 인수
    training_args = TrainingArguments()

    trainer = Trainer(model=model, args=training_args)

    # 체크포인트 디렉터리를 사용하여 체크포인트에서 트레이닝을 재개해야 합니다.
    trainer.train(resume_from_checkpoint=checkpoint_dir)

# Trainer를 정상적으로 인스턴스화
trainer = Trainer()

# Trainer 오브젝트를 전달하여 새로운 로깅 콜백 인스턴스화
evals_callback = WandbEvalsCallback(trainer, tokenizer, ...)

# Trainer에 콜백 추가
trainer.add_callback(evals_callback)

# Trainer 트레이닝을 정상적으로 시작
trainer.train()

from transformers.integrations import WandbCallback
import pandas as pd


def decode_predictions(tokenizer, predictions):
    labels = tokenizer.batch_decode(predictions.label_ids)
    logits = predictions.predictions.argmax(axis=-1)
    prediction_text = tokenizer.batch_decode(logits)
    return {"labels": labels, "predictions": prediction_text}


class WandbPredictionProgressCallback(WandbCallback):
    """트레이닝 중에 모델 예측을 로깅하는 사용자 지정 WandbCallback입니다.

    이 콜백은 트레이닝 중 각 로깅 단계에서 모델 예측과 레이블을 wandb.Table에 로깅합니다.
    트레이닝이 진행됨에 따라 모델 예측을 시각화할 수 있습니다.

    특성:
        trainer (Trainer): Hugging Face Trainer 인스턴스입니다.
        tokenizer (AutoTokenizer): 모델과 연결된 토크나이저입니다.
        sample_dataset (Dataset): 예측 생성을 위한 유효성 검사 데이터셋의 서브셋입니다.
        num_samples (int, optional): 예측 생성을 위해 유효성 검사 데이터셋에서 선택할 샘플 수입니다. 기본값은 100입니다.
        freq (int, optional): 로깅 빈도입니다. 기본값은 2입니다.
    """

    def __init__(self, trainer, tokenizer, val_dataset, num_samples=100, freq=2):
        """WandbPredictionProgressCallback 인스턴스를 초기화합니다.

        인수:
            trainer (Trainer): Hugging Face Trainer 인스턴스입니다.
            tokenizer (AutoTokenizer): 모델과 연결된 토크나이저입니다.
            val_dataset (Dataset): 유효성 검사 데이터셋입니다.
            num_samples (int, optional): 예측 생성을 위해 유효성 검사 데이터셋에서 선택할 샘플 수입니다.
              기본값은 100입니다.
            freq (int, optional): 로깅 빈도입니다. 기본값은 2입니다.
        """
        super().__init__()
        self.trainer = trainer
        self.tokenizer = tokenizer
        self.sample_dataset = val_dataset.select(range(num_samples))
        self.freq = freq

    def on_evaluate(self, args, state, control, **kwargs):
        super().on_evaluate(args, state, control, **kwargs)
        # `freq` 에포크마다 예측을 로깅하여 로깅 빈도 제어
        if state.epoch % self.freq == 0:
            # 예측 생성
            predictions = self.trainer.predict(self.sample_dataset)
            # 예측 및 레이블 디코딩
            predictions = decode_predictions(self.tokenizer, predictions)
            # wandb.Table에 예측 추가
            predictions_df = pd.DataFrame(predictions)
            predictions_df["epoch"] = state.epoch
            records_table = self._wandb.Table(dataframe=predictions_df)
            # 테이블을 wandb에 로깅
            self._wandb.log({"sample_predictions": records_table})


# 먼저 Trainer 인스턴스화
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=lm_datasets["train"],
    eval_dataset=lm_datasets["validation"],
)

# WandbPredictionProgressCallback 인스턴스화
progress_callback = WandbPredictionProgressCallback(
    trainer=trainer,
    tokenizer=tokenizer,
    val_dataset=lm_dataset["validation"],
    num_samples=10,
    freq=2,
)

# Trainer에 콜백 추가
trainer.add_callback(progress_callback)

wandb.init(
    project="amazon_sentiment_analysis",
    name="bert-base-high-lr",
    tags=["baseline", "high-lr"],
    group="bert",
)