Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

1
import argparse
2
import logging
3
import os
4
import random
5
import sys
6

7
import numpy as np
8
import torch
9
from datasets import load_from_disk, load_metric
10
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
11
from transformers.trainer_utils import get_last_checkpoint
12

13
if __name__ == "__main__":
14

15
    parser = argparse.ArgumentParser()
16

17
    # hyperparameters sent by the client are passed as command-line arguments to the script.
18
    parser.add_argument("--epochs", type=int, default=3)
19
    parser.add_argument("--train_batch_size", type=int, default=32)
20
    parser.add_argument("--eval_batch_size", type=int, default=64)
21
    parser.add_argument("--warmup_steps", type=int, default=500)
22
    parser.add_argument("--model_id", type=str)
23
    parser.add_argument("--learning_rate", type=str, default=5e-5)
24
    parser.add_argument("--fp16", type=bool, default=True)
25

26
    # Push to Hub Parameters
27
    parser.add_argument("--push_to_hub", type=bool, default=True)
28
    parser.add_argument("--hub_model_id", type=str, default=None)
29
    parser.add_argument("--hub_strategy", type=str, default=None)
30
    parser.add_argument("--hub_token", type=str, default=None)
31

32
    # Data, model, and output directories
33
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
34
    parser.add_argument("--output_dir", type=str, default=os.environ["SM_MODEL_DIR"])
35
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
36
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
37
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])
38

39
    args, _ = parser.parse_known_args()
40

41
    # make sure we have required parameters to push
42
    if args.push_to_hub:
43
        if args.hub_strategy is None:
44
            raise ValueError("--hub_strategy is required when pushing to Hub")
45
        if args.hub_token is None:
46
            raise ValueError("--hub_token is required when pushing to Hub")
47

48
    # sets hub id if not provided
49
    if args.hub_model_id is None:
50
        args.hub_model_id = args.model_id.replace("/", "--")
51

52
    # Set up logging
53
    logger = logging.getLogger(__name__)
54

55
    logging.basicConfig(
56
        level=logging.getLevelName("INFO"),
57
        handlers=[logging.StreamHandler(sys.stdout)],
58
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
59
    )
60

61
    # load datasets
62
    train_dataset = load_from_disk(args.training_dir)
63
    test_dataset = load_from_disk(args.test_dir)
64

65
    logger.info(f" loaded train_dataset length is: {len(train_dataset)}")
66
    logger.info(f" loaded test_dataset length is: {len(test_dataset)}")
67

68
    # define metrics and metrics function
69
    metric = load_metric("accuracy")
70

71
    def compute_metrics(eval_pred):
72
        predictions, labels = eval_pred
73
        predictions = np.argmax(predictions, axis=1)
74
        return metric.compute(predictions=predictions, references=labels)
75

76
    # Prepare model labels - useful in inference API
77
    labels = train_dataset.features["labels"].names
78
    num_labels = len(labels)
79
    label2id, id2label = dict(), dict()
80
    for i, label in enumerate(labels):
81
        label2id[label] = str(i)
82
        id2label[str(i)] = label
83

84
    # download model from model hub
85
    model = AutoModelForSequenceClassification.from_pretrained(
86
        args.model_id, num_labels=num_labels, label2id=label2id, id2label=id2label
87
    )
88
    tokenizer = AutoTokenizer.from_pretrained(args.model_id)
89

90
    # define training args
91
    training_args = TrainingArguments(
92
        output_dir=args.output_dir,
93
        overwrite_output_dir=True if get_last_checkpoint(args.output_dir) is not None else False,
94
        num_train_epochs=args.epochs,
95
        per_device_train_batch_size=args.train_batch_size,
96
        per_device_eval_batch_size=args.eval_batch_size,
97
        warmup_steps=args.warmup_steps,
98
        fp16=args.fp16,
99
        evaluation_strategy="epoch",
100
        save_strategy="epoch",
101
        save_total_limit=2,
102
        logging_dir=f"{args.output_data_dir}/logs",
103
        learning_rate=float(args.learning_rate),
104
        load_best_model_at_end=True,
105
        metric_for_best_model="accuracy",
106
        # push to hub parameters
107
        push_to_hub=args.push_to_hub,
108
        hub_strategy=args.hub_strategy,
109
        hub_model_id=args.hub_model_id,
110
        hub_token=args.hub_token,
111
    )
112

113
    # create Trainer instance
114
    trainer = Trainer(
115
        model=model,
116
        args=training_args,
117
        compute_metrics=compute_metrics,
118
        train_dataset=train_dataset,
119
        eval_dataset=test_dataset,
120
        tokenizer=tokenizer,
121
    )
122

123
    # train model
124
    trainer.train()
125

126
    # evaluate model
127
    eval_result = trainer.evaluate(eval_dataset=test_dataset)
128

129
    # save best model, metrics and create model card
130
    trainer.create_model_card(model_name=args.hub_model_id)
131
    trainer.push_to_hub()
132

133
    # Saves the model to s3 uses os.environ["SM_MODEL_DIR"] to make sure checkpointing works
134
    trainer.save_model(os.environ["SM_MODEL_DIR"])
135

136