Hugging Face Transformers

os.environ["WANDB_PROJECT"] = "<my-amazing-project>"  # W&B project に名前を付ける
os.environ["WANDB_LOG_MODEL"] = "checkpoint"  # すべてのモデル チェックポイントをログする

from transformers import TrainingArguments, Trainer

args = TrainingArguments(..., report_to="wandb")  # W&B ログを有効にする
trainer = Trainer(..., args=args)

run = wandb.init()
trainer.train()  # トレーニングを開始し、W&B にログする

# トレーニング後の分析、テスト、その他のログ対象コード

run.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 から wandb の run を再開する
with wandb.init(
    project=os.environ["WANDB_PROJECT"],
    id=last_run_id,
    resume="must",
) as run:
    # run に Artifact を接続する
    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 にログします。トレーニングの進行に伴い、モデルの予測を
    可視化することができます。

    Attributes:
        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 インスタンスを初期化します。

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