Add a transformers example (#322)

Author: ParagEkbote

.github/workflows/transformers.yml

@@ -0,0 +1,38 @@
+name: transformers
+
+on:
+  schedule:
+    - cron: '0 15 * * *'
+  pull_request:
+    paths:
+      - 'transformers/**'
+      - '.github/workflows/transformers.yml'
+  workflow_dispatch:
+
+jobs:
+  examples:
+    if: (github.event_name == 'schedule' && github.repository == 'optuna/optuna-examples') || (github.event_name != 'schedule')
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.10', '3.11', '3.12']
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: setup-python${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install (Python)
+      run: |
+        python -m pip install --upgrade pip
+        pip install --progress-bar off -U setuptools
+        pip install git+https://github.com/optuna/optuna.git
+        python -c 'import optuna'
+
+        pip install -r transformers/requirements.txt
+    - name: Run examples
+      run: |
+        python transformers/transformers_simple.py
+      env:
+        OMP_NUM_THREADS: 1

transformers/requirements.txt

@@ -0,0 +1,5 @@
+accelerate
+datasets
+evaluate
+scikit-learn
+transformers

transformers/transformers_simple.py

@@ -0,0 +1,96 @@
+"""
+Optuna example for fine-tuning a BERT-based text classification model on the IMDb dataset
+with hyperparameter optimization using Optuna. In this example, we fine-tune a lightweight
+pre-trained BERT model on a small subset of the IMDb dataset to classify movie reviews as
+positive or negative. We optimize the validation accuracy by tuning the learning rate
+and batch size. To learn more about transformers' hyperparameter search,
+you can check the following documentation:
+https://huggingface.co/docs/transformers/en/hpo_train.
+"""
+
+from datasets import load_dataset
+import evaluate
+
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+from transformers import set_seed
+from transformers import Trainer
+from transformers import TrainingArguments
+
+
+set_seed(42)
+
+
+train_dataset = load_dataset("imdb", split="train").shuffle(seed=42).select(range(1000))
+valid_dataset = load_dataset("imdb", split="test").shuffle(seed=42).select(range(500))
+
+model_name = "prajjwal1/bert-tiny"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+
+def tokenize(batch):
+    return tokenizer(batch["text"], padding="max_length", truncation=True, max_length=512)
+
+
+tokenized_train = train_dataset.map(tokenize, batched=True).select_columns(
+    ["input_ids", "attention_mask", "label"]
+)
+tokenized_valid = valid_dataset.map(tokenize, batched=True).select_columns(
+    ["input_ids", "attention_mask", "label"]
+)
+
+
+metric = evaluate.load("accuracy")
+
+
+def model_init():
+    return AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
+
+
+def compute_metrics(eval_pred):
+    predictions = eval_pred.predictions.argmax(axis=-1)
+    labels = eval_pred.label_ids
+    return metric.compute(predictions=predictions, references=labels)
+
+
+def compute_objective(metrics):
+    return metrics["eval_accuracy"]
+
+
+training_args = TrainingArguments(
+    eval_strategy="epoch",
+    save_strategy="best",
+    load_best_model_at_end=True,
+    logging_strategy="epoch",
+    report_to="none",
+)
+
+
+trainer = Trainer(
+    model_init=model_init,
+    args=training_args,
+    train_dataset=tokenized_train,
+    eval_dataset=tokenized_valid,
+    processing_class=tokenizer,
+    compute_metrics=compute_metrics,
+)
+
+
+def optuna_hp_space(trial):
+    return {
+        "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
+        "per_device_train_batch_size": trial.suggest_categorical(
+            "per_device_train_batch_size", [16, 32, 64, 128]
+        ),
+    }
+
+
+best_run = trainer.hyperparameter_search(
+    direction="maximize",
+    backend="optuna",
+    hp_space=optuna_hp_space,
+    n_trials=5,
+    compute_objective=compute_objective,
+)
+
+print(best_run)