run_tf_glue.py 2.59 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
import tensorflow as tf
import tensorflow_datasets
3
from transformers import *
thomwolf's avatar
thomwolf committed
4

thomwolf's avatar
thomwolf committed
5
# Load dataset, tokenizer, model from pretrained model/vocabulary
thomwolf's avatar
thomwolf committed
6
tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
thomwolf's avatar
thomwolf committed
7
model = TFBertForSequenceClassification.from_pretrained('bert-base-cased')
thomwolf's avatar
thomwolf committed
8
data = tensorflow_datasets.load('glue/mrpc')
thomwolf's avatar
thomwolf committed
9

thomwolf's avatar
thomwolf committed
10
# Prepare dataset for GLUE as a tf.data.Dataset instance
thomwolf's avatar
thomwolf committed
11
12
13
train_dataset = glue_convert_examples_to_features(data['train'], tokenizer, 128, 'mrpc')
valid_dataset = glue_convert_examples_to_features(data['validation'], tokenizer, 128, 'mrpc')
train_dataset = train_dataset.shuffle(100).batch(32).repeat(2)
thomwolf's avatar
thomwolf committed
14
15
valid_dataset = valid_dataset.batch(64)

thomwolf's avatar
thomwolf committed
16
# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule 
thomwolf's avatar
thomwolf committed
17
optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
thomwolf's avatar
thomwolf committed
18
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
thomwolf's avatar
thomwolf committed
19
20
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
model.compile(optimizer=optimizer, loss=loss, metrics=[metric])
thomwolf's avatar
thomwolf committed
21
22

# Train and evaluate using tf.keras.Model.fit()
thomwolf's avatar
thomwolf committed
23
24
25
26
27
28
29
30
history = model.fit(train_dataset, epochs=2, steps_per_epoch=115,
                    validation_data=valid_dataset, validation_steps=7)

>>> Train for 115 steps, validate for 7 steps
>>> Epoch 1/2
>>> 115/115 [==============================] - 53s 459ms/step - loss: 0.6033 - accuracy: 0.6712 - val_loss: 0.4964 - val_accuracy: 0.7647
>>> Epoch 2/2
>>> 115/115 [==============================] - 33s 289ms/step - loss: 0.4141 - accuracy: 0.8160 - val_loss: 0.3914 - val_accuracy: 0.8382
thomwolf's avatar
thomwolf committed
31

thomwolf's avatar
thomwolf committed
32
# Load the TensorFlow model in PyTorch for inspection
thomwolf's avatar
thomwolf committed
33
34
model.save_pretrained('./save/')
pytorch_model = BertForSequenceClassification.from_pretrained('./save/', from_tf=True)
thomwolf's avatar
thomwolf committed
35

thomwolf's avatar
thomwolf committed
36
37
38
39
40
41
42
43
44
45
46
47
48
# Quickly test a few predictions - MRPC is a paraphrasing task, let's see if our model learned the task
sentence_0 = "This research was consistent with his findings."
sentence_1 = "His findings were compatible with this research."
sentence_2 = "His findings were not compatible with this research."
inputs_1 = tokenizer.encode_plus(sentence_0, sentence_1, add_special_tokens=True, return_tensors='pt')
inputs_2 = tokenizer.encode_plus(sentence_0, sentence_2, add_special_tokens=True, return_tensors='pt')

pred_1 = pytorch_model(**inputs_1)[0].argmax().item()
pred_2 = pytorch_model(**inputs_2)[0].argmax().item()
print("sentence_1 is", "a paraphrase" if pred_1 else "not a paraphrase", "of sentence_0")
print("sentence_2 is", "a paraphrase" if pred_2 else "not a paraphrase", "of sentence_0")
>>> sentence_1 is a paraphrase of sentence_0
>>> sentence_2 is not a paraphrase of sentence_0