Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ffc522e1a1 | ||
|
|
459c7707f0 | ||
|
|
97debe4723 | ||
|
|
9a9e68e7b5 | ||
|
|
74577316b5 | ||
|
|
2b6f10e4e3 | ||
|
|
59988ba723 |
+311
-87
@@ -1,3 +1,22 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2018 The HuggingFace Inc. team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" Finetuning the library models for sequence classification on GLUE (Bert, XLM, XLNet, RoBERTa)."""
|
||||
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
|
||||
import tensorflow as tf
|
||||
@@ -5,104 +24,309 @@ import tensorflow_datasets
|
||||
|
||||
from transformers import (
|
||||
BertConfig,
|
||||
BertForSequenceClassification,
|
||||
BertTokenizer,
|
||||
DistilBertConfig,
|
||||
DistilBertTokenizer,
|
||||
RobertaConfig,
|
||||
RobertaTokenizer,
|
||||
TFBertForSequenceClassification,
|
||||
glue_convert_examples_to_features,
|
||||
glue_processors,
|
||||
TFDistilBertForSequenceClassification,
|
||||
TFRobertaForSequenceClassification,
|
||||
TFXLMForSequenceClassification,
|
||||
TFXLNetForSequenceClassification,
|
||||
XLMConfig,
|
||||
XLMTokenizer,
|
||||
XLNetConfig,
|
||||
XLNetTokenizer,
|
||||
)
|
||||
from transformers import glue_convert_examples_to_features as convert_examples_to_features
|
||||
from transformers import glue_output_modes as output_modes
|
||||
from transformers import glue_processors as processors
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ALL_MODELS = sum(
|
||||
(
|
||||
tuple(conf.pretrained_config_archive_map.keys())
|
||||
for conf in (BertConfig, XLNetConfig, XLMConfig, RobertaConfig, DistilBertConfig)
|
||||
),
|
||||
(),
|
||||
)
|
||||
|
||||
|
||||
# script parameters
|
||||
BATCH_SIZE = 32
|
||||
EVAL_BATCH_SIZE = BATCH_SIZE * 2
|
||||
USE_XLA = False
|
||||
USE_AMP = False
|
||||
EPOCHS = 3
|
||||
|
||||
TASK = "mrpc"
|
||||
|
||||
if TASK == "sst-2":
|
||||
TFDS_TASK = "sst2"
|
||||
elif TASK == "sts-b":
|
||||
TFDS_TASK = "stsb"
|
||||
else:
|
||||
TFDS_TASK = TASK
|
||||
|
||||
num_labels = len(glue_processors[TASK]().get_labels())
|
||||
print(num_labels)
|
||||
|
||||
tf.config.optimizer.set_jit(USE_XLA)
|
||||
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": USE_AMP})
|
||||
|
||||
# Load tokenizer and model from pretrained model/vocabulary. Specify the number of labels to classify (2+: classification, 1: regression)
|
||||
config = BertConfig.from_pretrained("bert-base-cased", num_labels=num_labels)
|
||||
tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
|
||||
model = TFBertForSequenceClassification.from_pretrained("bert-base-cased", config=config)
|
||||
|
||||
# Load dataset via TensorFlow Datasets
|
||||
data, info = tensorflow_datasets.load(f"glue/{TFDS_TASK}", with_info=True)
|
||||
train_examples = info.splits["train"].num_examples
|
||||
|
||||
# MNLI expects either validation_matched or validation_mismatched
|
||||
valid_examples = info.splits["validation"].num_examples
|
||||
|
||||
# Prepare dataset for GLUE as a tf.data.Dataset instance
|
||||
train_dataset = glue_convert_examples_to_features(data["train"], tokenizer, 128, TASK)
|
||||
|
||||
# MNLI expects either validation_matched or validation_mismatched
|
||||
valid_dataset = glue_convert_examples_to_features(data["validation"], tokenizer, 128, TASK)
|
||||
train_dataset = train_dataset.shuffle(128).batch(BATCH_SIZE).repeat(-1)
|
||||
valid_dataset = valid_dataset.batch(EVAL_BATCH_SIZE)
|
||||
|
||||
# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule
|
||||
opt = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
|
||||
if USE_AMP:
|
||||
# loss scaling is currently required when using mixed precision
|
||||
opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, "dynamic")
|
||||
MODEL_CLASSES = {
|
||||
"bert": (BertConfig, TFBertForSequenceClassification, BertTokenizer),
|
||||
"xlnet": (XLNetConfig, TFXLNetForSequenceClassification, XLNetTokenizer),
|
||||
"xlm": (XLMConfig, TFXLMForSequenceClassification, XLMTokenizer),
|
||||
"roberta": (RobertaConfig, TFRobertaForSequenceClassification, RobertaTokenizer),
|
||||
"distilbert": (DistilBertConfig, TFDistilBertForSequenceClassification, DistilBertTokenizer),
|
||||
}
|
||||
|
||||
|
||||
if num_labels == 1:
|
||||
loss = tf.keras.losses.MeanSquaredError()
|
||||
else:
|
||||
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
|
||||
def load_and_cache_examples(args, data, task, tokenizer, split):
|
||||
if task == "mnli" and split == "validation":
|
||||
split = "validation_matched"
|
||||
|
||||
metric = tf.keras.metrics.SparseCategoricalAccuracy("accuracy")
|
||||
model.compile(optimizer=opt, loss=loss, metrics=[metric])
|
||||
features_output_dir = os.path.join(args.output_dir, "features")
|
||||
cached_features_file = os.path.join(
|
||||
features_output_dir,
|
||||
"cached_{}_{}_{}_{}.tfrecord".format(
|
||||
split, list(filter(None, args.model_name_or_path.split("/"))).pop(), str(args.max_seq_length), str(task)
|
||||
),
|
||||
)
|
||||
|
||||
# Train and evaluate using tf.keras.Model.fit()
|
||||
train_steps = train_examples // BATCH_SIZE
|
||||
valid_steps = valid_examples // EVAL_BATCH_SIZE
|
||||
if not os.path.exists(cached_features_file) or args.overwrite_cache:
|
||||
logger.info("Converting examples to features")
|
||||
dataset = convert_examples_to_features(data[split], tokenizer, args.max_seq_length, task)
|
||||
|
||||
history = model.fit(
|
||||
train_dataset,
|
||||
epochs=EPOCHS,
|
||||
steps_per_epoch=train_steps,
|
||||
validation_data=valid_dataset,
|
||||
validation_steps=valid_steps,
|
||||
)
|
||||
if not os.path.exists(features_output_dir):
|
||||
os.makedirs(features_output_dir)
|
||||
|
||||
# Save TF2 model
|
||||
os.makedirs("./save/", exist_ok=True)
|
||||
model.save_pretrained("./save/")
|
||||
with tf.compat.v1.python_io.TFRecordWriter(cached_features_file) as tfwriter:
|
||||
for feature in dataset:
|
||||
example, label = feature
|
||||
feature_key_value_pair = {
|
||||
"input_ids": tf.train.Feature(int64_list=tf.train.Int64List(value=example["input_ids"])),
|
||||
"attention_mask": tf.train.Feature(int64_list=tf.train.Int64List(value=example["attention_mask"])),
|
||||
"token_type_ids": tf.train.Feature(int64_list=tf.train.Int64List(value=example["token_type_ids"])),
|
||||
"label": tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
|
||||
}
|
||||
features = tf.train.Features(feature=feature_key_value_pair)
|
||||
example = tf.train.Example(features=features)
|
||||
|
||||
if TASK == "mrpc":
|
||||
# Load the TensorFlow model in PyTorch for inspection
|
||||
# This is to demo the interoperability between the two frameworks, you don't have to
|
||||
# do this in real life (you can run the inference on the TF model).
|
||||
pytorch_model = BertForSequenceClassification.from_pretrained("./save/", from_tf=True)
|
||||
tfwriter.write(example.SerializeToString())
|
||||
|
||||
# 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")
|
||||
logger.info("Features saved to cache")
|
||||
|
||||
del inputs_1["special_tokens_mask"]
|
||||
del inputs_2["special_tokens_mask"]
|
||||
features = {
|
||||
"input_ids": tf.io.FixedLenFeature([args.max_seq_length], tf.int64),
|
||||
"attention_mask": tf.io.FixedLenFeature([args.max_seq_length], tf.int64),
|
||||
"token_type_ids": tf.io.FixedLenFeature([args.max_seq_length], tf.int64),
|
||||
"label": tf.io.FixedLenFeature([], tf.int64),
|
||||
}
|
||||
|
||||
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")
|
||||
def select_data_from_record(record):
|
||||
record = tf.io.parse_single_example(record, features)
|
||||
x = {
|
||||
"input_ids": record["input_ids"],
|
||||
"attention_mask": record["attention_mask"],
|
||||
"token_type_ids": record["token_type_ids"],
|
||||
}
|
||||
y = record["label"]
|
||||
return (x, y)
|
||||
|
||||
dataset = tf.data.TFRecordDataset(cached_features_file)
|
||||
dataset = dataset.map(select_data_from_record)
|
||||
|
||||
logger.info("Created dataset %s from TFRecord" % split)
|
||||
return dataset
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument(
|
||||
"--model_type",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_name_or_path",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(ALL_MODELS),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--task_name",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help="The name of the task to train selected in the list: " + ", ".join(processors.keys()),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help="The output directory where the model predictions and checkpoints will be written.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite_output_dir", action="store_true", help="Overwrite the content of the output directory"
|
||||
)
|
||||
|
||||
parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
|
||||
parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")
|
||||
parser.add_argument(
|
||||
"--evaluate_during_training", action="store_true", help="Rul evaluation during training at each logging step."
|
||||
)
|
||||
|
||||
parser.add_argument("--train_batch_size", default=8, type=int, help="Batch size per GPU/CPU for training.")
|
||||
parser.add_argument(
|
||||
"--valid_batch_size", default=8, type=int, help="Batch size per GPU/CPU for validation during training."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--test_batch_size", default=8, type=int, help="Batch size per GPU/CPU for evaluation after training."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
|
||||
)
|
||||
|
||||
parser.add_argument("--num_train_epochs", default=3, type=int, help="Total number of training epochs to perform.")
|
||||
parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
|
||||
parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
|
||||
parser.add_argument(
|
||||
"--max_seq_length",
|
||||
default=128,
|
||||
type=int,
|
||||
help="The maximum total input sequence length after tokenization. Sequences longer "
|
||||
"than this will be truncated, sequences shorter will be padded.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--do_lower_case", action="store_true", help="Set this flag if you are using an uncased model."
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tokenizer_name",
|
||||
default="",
|
||||
type=str,
|
||||
help="Pretrained tokenizer name or path if not the same as model_name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache_dir",
|
||||
default="",
|
||||
type=str,
|
||||
help="Where do you want to store the pre-trained models downloaded from s3",
|
||||
)
|
||||
parser.add_argument("--xla", action="store_true", help="Whether to use XLA (Accelerated Linear Algebra).")
|
||||
parser.add_argument("--amp", action="store_true", help="Whether to use AMP (Automatic Mixed Precision).")
|
||||
parser.add_argument(
|
||||
"--force_download",
|
||||
action="store_true",
|
||||
help="Whether to force download the weights from S3 (useful if the file is corrupted).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
)
|
||||
|
||||
if os.path.exists(args.output_dir) and args.do_train:
|
||||
if not args.overwrite_output_dir and bool(
|
||||
[file for file in os.listdir(args.output_dir) if "features" not in file]
|
||||
):
|
||||
raise ValueError(
|
||||
"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
|
||||
args.output_dir
|
||||
)
|
||||
)
|
||||
|
||||
args.model_type = args.model_type.lower()
|
||||
config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
|
||||
|
||||
TASK = args.task_name.lower()
|
||||
|
||||
if TASK not in processors:
|
||||
raise ValueError("Task not found: %s" % (TASK))
|
||||
|
||||
if TASK == "sst-2":
|
||||
TFDS_TASK = "sst2"
|
||||
elif TASK == "sts-b":
|
||||
TFDS_TASK = "stsb"
|
||||
else:
|
||||
TFDS_TASK = TASK
|
||||
|
||||
num_labels = len(processors[TASK]().get_labels())
|
||||
print(num_labels)
|
||||
|
||||
tf.config.optimizer.set_jit(args.xla)
|
||||
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.amp})
|
||||
|
||||
# Load tokenizer and model from pretrained model/vocabulary. Specify the number of labels to classify (2+: classification, 1: regression)
|
||||
config = config_class.from_pretrained(
|
||||
args.config_name if args.config_name else args.model_name_or_path,
|
||||
num_labels=num_labels,
|
||||
finetuning_task=args.task_name,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
force_download=args.force_download,
|
||||
)
|
||||
|
||||
tokenizer = tokenizer_class.from_pretrained(
|
||||
args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
|
||||
do_lower_case=args.do_lower_case,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
force_download=args.force_download,
|
||||
)
|
||||
|
||||
model = model_class.from_pretrained(
|
||||
args.model_name_or_path,
|
||||
config=config,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
force_download=args.force_download,
|
||||
)
|
||||
|
||||
# Load dataset via TensorFlow Datasets
|
||||
data, info = tensorflow_datasets.load("glue/%s" % TFDS_TASK, with_info=True)
|
||||
|
||||
# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule
|
||||
opt = tf.keras.optimizers.Adam(learning_rate=args.learning_rate, epsilon=args.adam_epsilon)
|
||||
|
||||
if args.amp:
|
||||
# loss scaling is currently required when using mixed precision
|
||||
opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, "dynamic")
|
||||
|
||||
if num_labels == 1:
|
||||
loss = tf.keras.losses.MeanSquaredError()
|
||||
else:
|
||||
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
|
||||
|
||||
metric = tf.keras.metrics.SparseCategoricalAccuracy("accuracy")
|
||||
model.compile(optimizer=opt, loss=loss, metrics=[metric])
|
||||
|
||||
class save_model(tf.keras.callbacks.Callback):
|
||||
def on_epoch_end(self, epoch, logs=None):
|
||||
print("Saving model at epoch {}".format(epoch))
|
||||
output_dir = os.path.join(args.output_dir, "checkpoint-epoch-{}".format(epoch))
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
self.model.save_pretrained(output_dir)
|
||||
|
||||
if args.do_train:
|
||||
train_dataset = load_and_cache_examples(args, data=data, task=TASK, tokenizer=tokenizer, split="train")
|
||||
train_dataset = train_dataset.batch(args.train_batch_size).repeat(args.num_train_epochs)
|
||||
train_examples = info.splits["train"].num_examples / args.train_batch_size
|
||||
|
||||
validation_identifier = "validation_mismatched" if TASK == "mnli" else "validation"
|
||||
valid_dataset = load_and_cache_examples(
|
||||
args, data=data, task=TASK, tokenizer=tokenizer, split=validation_identifier
|
||||
)
|
||||
valid_dataset = valid_dataset.batch(args.valid_batch_size)
|
||||
valid_examples = info.splits[validation_identifier].num_examples / args.valid_batch_size
|
||||
|
||||
history = model.fit(
|
||||
train_dataset,
|
||||
steps_per_epoch=train_examples,
|
||||
epochs=args.num_train_epochs,
|
||||
validation_data=valid_dataset if args.evaluate_during_training else None,
|
||||
validation_steps=valid_examples if args.evaluate_during_training else None,
|
||||
callbacks=[save_model()],
|
||||
)
|
||||
|
||||
if args.do_eval:
|
||||
test_dataset = load_and_cache_examples(args, data=data, task=TASK, tokenizer=tokenizer, split="test")
|
||||
test_dataset = test_dataset.batch(args.test_batch_size)
|
||||
test_examples = info.splits["test"].num_examples / args.test_batch_size
|
||||
|
||||
results = model.evaluate(test_dataset, steps=test_examples)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,699 @@
|
||||
# Copyright 2018 The HuggingFace Inc. team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" Finetuning the library models for sequence classification on GLUE (Bert, XLM, XLNet, RoBERTa)."""
|
||||
from __future__ import absolute_import, division, print_function
|
||||
|
||||
import argparse
|
||||
import datetime
|
||||
import glob
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import regex as re
|
||||
import tensorflow as tf
|
||||
from fastprogress import master_bar, progress_bar
|
||||
from seqeval import metrics
|
||||
from tqdm import tqdm, trange
|
||||
|
||||
from transformers import (
|
||||
TF2_WEIGHTS_NAME,
|
||||
BertConfig,
|
||||
BertTokenizer,
|
||||
DistilBertConfig,
|
||||
DistilBertTokenizer,
|
||||
GradientAccumulator,
|
||||
SquadV1Processor,
|
||||
SquadV2Processor,
|
||||
TFBertForQuestionAnswering,
|
||||
TFDistilBertForQuestionAnswering,
|
||||
TFXLMForQuestionAnsweringSimple,
|
||||
TFXLNetForQuestionAnsweringSimple,
|
||||
XLMConfig,
|
||||
XLMTokenizer,
|
||||
XLNetConfig,
|
||||
XLNetTokenizer,
|
||||
create_optimizer,
|
||||
squad_convert_examples_to_features,
|
||||
)
|
||||
from transformers.data.metrics.squad_metrics import compute_predictions_logits, squad_evaluate
|
||||
from transformers.data.processors.squad import SquadResult
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ALL_MODELS = sum(
|
||||
(
|
||||
tuple(conf.pretrained_config_archive_map.keys())
|
||||
for conf in (BertConfig, XLNetConfig, XLMConfig, DistilBertConfig)
|
||||
),
|
||||
(),
|
||||
)
|
||||
|
||||
MODEL_CLASSES = {
|
||||
"bert": (BertConfig, TFBertForQuestionAnswering, BertTokenizer),
|
||||
"xlnet": (XLNetConfig, TFXLNetForQuestionAnsweringSimple, XLNetTokenizer),
|
||||
"xlm": (XLMConfig, TFXLMForQuestionAnsweringSimple, XLMTokenizer),
|
||||
"distilbert": (DistilBertConfig, TFDistilBertForQuestionAnswering, DistilBertTokenizer),
|
||||
}
|
||||
|
||||
|
||||
def train(args, strategy, train_dataset, tokenizer, model, num_train_examples, train_batch_size):
|
||||
if args.max_steps > 0:
|
||||
num_train_steps = args.max_steps * args.gradient_accumulation_steps
|
||||
args.num_train_epochs = 1
|
||||
else:
|
||||
num_train_steps = (
|
||||
math.ceil(num_train_examples / train_batch_size)
|
||||
// args.gradient_accumulation_steps
|
||||
* args.num_train_epochs
|
||||
)
|
||||
|
||||
writer = tf.summary.create_file_writer("/tmp/mylogs")
|
||||
|
||||
with strategy.scope():
|
||||
loss_fct = tf.keras.losses.SparseCategoricalCrossentropy(
|
||||
reduction=tf.keras.losses.Reduction.NONE, from_logits=True
|
||||
)
|
||||
optimizer = create_optimizer(args.learning_rate, num_train_steps, args.warmup_steps)
|
||||
|
||||
if args.xla:
|
||||
tf.config.optimizer.set_jit(True)
|
||||
|
||||
if args.amp:
|
||||
optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(optimizer, "dynamic")
|
||||
|
||||
loss_metric = tf.keras.metrics.Mean(name="loss", dtype=tf.float32)
|
||||
gradient_accumulator = GradientAccumulator()
|
||||
|
||||
logging.info("***** Running training *****")
|
||||
logging.info(" Num examples = %d", num_train_examples)
|
||||
logging.info(" Num Epochs = %d", args.num_train_epochs)
|
||||
logging.info(" Instantaneous batch size per device = %d", args.per_device_train_batch_size)
|
||||
logging.info(
|
||||
" Total train batch size (w. parallel, distributed & accumulation) = %d",
|
||||
train_batch_size * args.gradient_accumulation_steps,
|
||||
)
|
||||
logging.info(" Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
|
||||
logging.info(" Total training steps = %d", num_train_steps)
|
||||
|
||||
model.summary()
|
||||
|
||||
@tf.function
|
||||
def apply_gradients():
|
||||
grads_and_vars = []
|
||||
|
||||
for gradient, variable in zip(gradient_accumulator.gradients, model.trainable_variables):
|
||||
if gradient is not None:
|
||||
scaled_gradient = gradient / (args.n_device * args.gradient_accumulation_steps)
|
||||
grads_and_vars.append((scaled_gradient, variable))
|
||||
else:
|
||||
grads_and_vars.append((gradient, variable))
|
||||
|
||||
optimizer.apply_gradients(grads_and_vars, args.max_grad_norm)
|
||||
gradient_accumulator.reset()
|
||||
|
||||
@tf.function
|
||||
def train_step(train_features, train_labels):
|
||||
def step_fn(train_features, train_labels):
|
||||
with tf.GradientTape() as tape:
|
||||
start_logits, end_logits = model(train_features)
|
||||
start_logits = tf.multiply(
|
||||
start_logits, tf.dtypes.cast((train_features["attention_mask"]), tf.float32)
|
||||
)
|
||||
end_logits = tf.multiply(end_logits, tf.dtypes.cast((train_features["attention_mask"]), tf.float32))
|
||||
start_loss = loss_fct(train_labels["start_position"], start_logits)
|
||||
end_loss = loss_fct(train_labels["end_position"], end_logits)
|
||||
total_loss = (start_loss + end_loss) / 2
|
||||
|
||||
loss = tf.reduce_sum(total_loss) * (1.0 / train_batch_size)
|
||||
grads = tape.gradient(loss, model.trainable_variables)
|
||||
|
||||
gradient_accumulator(grads)
|
||||
|
||||
return total_loss
|
||||
|
||||
per_example_losses = strategy.experimental_run_v2(step_fn, args=(train_features, train_labels))
|
||||
mean_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_example_losses, axis=0)
|
||||
|
||||
return mean_loss
|
||||
|
||||
current_time = datetime.datetime.now()
|
||||
train_iterator = master_bar(range(args.num_train_epochs))
|
||||
global_step = 0
|
||||
logging_loss = 0.0
|
||||
|
||||
for epoch in train_iterator:
|
||||
epoch_iterator = progress_bar(
|
||||
train_dataset,
|
||||
total=num_train_steps / args.num_train_epochs,
|
||||
parent=train_iterator,
|
||||
display=args.n_device > 1,
|
||||
)
|
||||
step = 1
|
||||
|
||||
with strategy.scope():
|
||||
for train_features, train_labels in tqdm(
|
||||
epoch_iterator,
|
||||
desc="Training, epoch {}".format(epoch),
|
||||
total=int(num_train_steps / args.num_train_epochs),
|
||||
):
|
||||
loss = train_step(train_features, train_labels)
|
||||
|
||||
if step % args.gradient_accumulation_steps == 0:
|
||||
strategy.experimental_run_v2(apply_gradients)
|
||||
|
||||
loss_metric(loss)
|
||||
|
||||
global_step += 1
|
||||
|
||||
if args.logging_steps > 0 and global_step % args.logging_steps == 0:
|
||||
# Log metrics
|
||||
if args.n_device == 1 and args.evaluate_during_training:
|
||||
# Only evaluate when single GPU otherwise metrics may not average well
|
||||
results = evaluate(args, strategy, model, tokenizer, global_step)
|
||||
|
||||
with writer.as_default():
|
||||
tf.summary.scalar("exact", results[0], global_step)
|
||||
tf.summary.scalar("f1", results[1], global_step)
|
||||
|
||||
lr = optimizer.learning_rate
|
||||
learning_rate = lr(step)
|
||||
|
||||
with writer.as_default():
|
||||
tf.summary.scalar("lr", learning_rate, global_step)
|
||||
tf.summary.scalar(
|
||||
"loss", (loss_metric.result() - logging_loss) / args.logging_steps, global_step
|
||||
)
|
||||
|
||||
logging_loss = loss_metric.result()
|
||||
|
||||
with writer.as_default():
|
||||
tf.summary.scalar("loss", loss_metric.result(), step=step)
|
||||
|
||||
if args.save_steps > 0 and global_step % args.save_steps == 0:
|
||||
# Save model checkpoint
|
||||
output_dir = os.path.join(args.output_dir, "checkpoint-{}".format(global_step))
|
||||
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
model.save_pretrained(output_dir)
|
||||
logging.info("Saving model checkpoint to %s", output_dir)
|
||||
|
||||
train_iterator.child.comment = f"loss : {loss_metric.result()}"
|
||||
step += 1
|
||||
|
||||
train_iterator.write(f"loss epoch {epoch + 1}: {loss_metric.result()}")
|
||||
|
||||
loss_metric.reset_states()
|
||||
|
||||
logging.info(" Training took time = {}".format(datetime.datetime.now() - current_time))
|
||||
|
||||
|
||||
def evaluate(args, strategy, model, tokenizer, prefix):
|
||||
eval_batch_size = args.per_device_eval_batch_size * args.n_device
|
||||
eval_dataset, size, examples, features = load_and_cache_examples(
|
||||
args, tokenizer, evaluate=True, output_examples=True
|
||||
)
|
||||
eval_dataset = eval_dataset.batch(args.per_device_eval_batch_size)
|
||||
eval_dataset = strategy.experimental_distribute_dataset(eval_dataset)
|
||||
num_eval_steps = math.ceil(size / eval_batch_size)
|
||||
master = master_bar(range(1))
|
||||
eval_iterator = progress_bar(eval_dataset, total=num_eval_steps, parent=master, display=args.n_device > 1)
|
||||
|
||||
logging.info("***** Running evaluation *****")
|
||||
logging.info(" Num examples = %d", size)
|
||||
logging.info(" Batch size = %d", eval_batch_size)
|
||||
|
||||
all_results = []
|
||||
for index, (eval_features, eval_labels) in tqdm(
|
||||
enumerate(eval_iterator), total=size / args.per_device_eval_batch_size
|
||||
):
|
||||
with strategy.scope():
|
||||
start_logits, end_logits = model(eval_features)
|
||||
|
||||
for sample_index in range(args.per_device_eval_batch_size):
|
||||
if index * args.per_device_eval_batch_size + sample_index < len(features):
|
||||
result = SquadResult(
|
||||
features[index * args.per_device_eval_batch_size + sample_index].unique_id,
|
||||
start_logits[sample_index],
|
||||
end_logits[sample_index],
|
||||
)
|
||||
all_results.append(result)
|
||||
|
||||
output_prediction_file = os.path.join(args.output_dir, "predictions_{}.json".format(prefix))
|
||||
output_nbest_file = os.path.join(args.output_dir, "nbest_predictions_{}.json".format(prefix))
|
||||
|
||||
if args.version_2_with_negative:
|
||||
output_null_log_odds_file = os.path.join(args.output_dir, "null_odds_{}.json".format(prefix))
|
||||
else:
|
||||
output_null_log_odds_file = None
|
||||
|
||||
predictions = compute_predictions_logits(
|
||||
examples,
|
||||
features,
|
||||
all_results,
|
||||
args.n_best_size,
|
||||
args.max_answer_length,
|
||||
args.do_lower_case,
|
||||
output_prediction_file,
|
||||
output_nbest_file,
|
||||
output_null_log_odds_file,
|
||||
args.verbose_logging,
|
||||
args.version_2_with_negative,
|
||||
args.null_score_diff_threshold,
|
||||
tokenizer,
|
||||
)
|
||||
|
||||
# Compute the F1 and exact scores.
|
||||
results = squad_evaluate(examples, predictions)
|
||||
return results
|
||||
|
||||
|
||||
def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False):
|
||||
features_output_dir = os.path.join(args.output_dir, "features")
|
||||
cached_features_file = os.path.join(
|
||||
features_output_dir,
|
||||
"cached_{}_{}_{}.tfrecord".format(
|
||||
"dev" if evaluate else "train",
|
||||
list(filter(None, args.model_name_or_path.split("/"))).pop(),
|
||||
str(args.max_seq_length),
|
||||
),
|
||||
)
|
||||
|
||||
if not os.path.exists(cached_features_file) or args.overwrite_cache:
|
||||
if args.version_2_with_negative:
|
||||
processor = SquadV2Processor()
|
||||
else:
|
||||
processor = SquadV1Processor()
|
||||
|
||||
if args.data_dir:
|
||||
directory = args.directory
|
||||
examples = processor.get_dev_examples(directory) if evaluate else processor.get_train_examples(directory)
|
||||
else:
|
||||
try:
|
||||
import tensorflow_datasets as tfds
|
||||
except ImportError:
|
||||
raise ImportError("If not data_dir is specified, tensorflow_datasets needs to be installed.")
|
||||
|
||||
if args.version_2_with_negative:
|
||||
logger.warning("tensorflow_datasets does not handle version 2 of SQuAD.")
|
||||
|
||||
examples = processor.get_examples_from_dataset(tfds.load("squad"), evaluate=evaluate)
|
||||
examples = examples
|
||||
|
||||
logger.info("Converting examples to features")
|
||||
features, dataset = squad_convert_examples_to_features(
|
||||
examples,
|
||||
tokenizer,
|
||||
args.max_seq_length,
|
||||
args.doc_stride,
|
||||
args.max_query_length,
|
||||
is_training=not evaluate,
|
||||
return_dataset="tf",
|
||||
)
|
||||
|
||||
if not os.path.exists(features_output_dir):
|
||||
os.makedirs(features_output_dir)
|
||||
|
||||
with tf.compat.v1.python_io.TFRecordWriter(cached_features_file) as tfwriter:
|
||||
for feature in tqdm(dataset, desc="Building tfrecord dataset", total=len(features)):
|
||||
example, result = feature
|
||||
feature_key_value_pair = {
|
||||
"input_ids": tf.train.Feature(int64_list=tf.train.Int64List(value=example["input_ids"])),
|
||||
"attention_mask": tf.train.Feature(int64_list=tf.train.Int64List(value=example["attention_mask"])),
|
||||
"token_type_ids": tf.train.Feature(int64_list=tf.train.Int64List(value=example["token_type_ids"])),
|
||||
"start_position": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[result["start_position"]])
|
||||
),
|
||||
"end_position": tf.train.Feature(int64_list=tf.train.Int64List(value=[result["end_position"]])),
|
||||
"cls_index": tf.train.Feature(int64_list=tf.train.Int64List(value=[result["cls_index"]])),
|
||||
"p_mask": tf.train.Feature(int64_list=tf.train.Int64List(value=result["p_mask"])),
|
||||
}
|
||||
feature_skeleton = tf.train.Features(feature=feature_key_value_pair)
|
||||
example = tf.train.Example(features=feature_skeleton)
|
||||
|
||||
tfwriter.write(example.SerializeToString())
|
||||
|
||||
with open("{}.pickle".format(cached_features_file.replace(".tfrecord", "_examples")), "wb") as handle:
|
||||
pickle.dump(examples, handle, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
with open("{}.pickle".format(cached_features_file.replace(".tfrecord", "_features")), "wb") as handle:
|
||||
pickle.dump(features, handle, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
logger.info("Features saved to cache")
|
||||
|
||||
feature_skeleton = {
|
||||
"input_ids": tf.io.FixedLenFeature([args.max_seq_length], tf.int64),
|
||||
"attention_mask": tf.io.FixedLenFeature([args.max_seq_length], tf.int64),
|
||||
"token_type_ids": tf.io.FixedLenFeature([args.max_seq_length], tf.int64),
|
||||
"start_position": tf.io.FixedLenFeature([], tf.int64),
|
||||
"end_position": tf.io.FixedLenFeature([], tf.int64),
|
||||
"cls_index": tf.io.FixedLenFeature([], tf.int64),
|
||||
"p_mask": tf.io.FixedLenFeature([args.max_seq_length], tf.int64),
|
||||
}
|
||||
|
||||
def select_data_from_record(record):
|
||||
record = tf.io.parse_single_example(record, feature_skeleton)
|
||||
x = {
|
||||
"input_ids": record["input_ids"],
|
||||
"attention_mask": record["attention_mask"],
|
||||
"token_type_ids": record["token_type_ids"],
|
||||
}
|
||||
y = {
|
||||
"start_position": record["start_position"],
|
||||
"end_position": record["end_position"],
|
||||
"cls_index": record["cls_index"],
|
||||
"p_mask": record["p_mask"],
|
||||
}
|
||||
return x, y
|
||||
|
||||
dataset = tf.data.TFRecordDataset(cached_features_file)
|
||||
dataset = dataset.map(select_data_from_record)
|
||||
|
||||
with open("{}.pickle".format(cached_features_file.replace(".tfrecord", "_examples")), "rb") as handle:
|
||||
examples = pickle.load(handle)
|
||||
|
||||
with open("{}.pickle".format(cached_features_file.replace(".tfrecord", "_features")), "rb") as handle:
|
||||
features = pickle.load(handle)
|
||||
|
||||
logger.info("Created dataset %s from TFRecord" % "dev" if evaluate else "train")
|
||||
|
||||
if output_examples:
|
||||
return dataset, len(list(dataset.__iter__())), examples, features
|
||||
return dataset, len(list(dataset.__iter__()))
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument(
|
||||
"--model_type",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_name_or_path",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(ALL_MODELS),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help="The output directory where the model predictions and checkpoints will be written.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data_dir",
|
||||
default=None,
|
||||
type=str,
|
||||
required=False,
|
||||
help="The input data directory containing the .json files. If no data dir is specified, uses tensorflow_datasets to load the data."
|
||||
+ ", ".join(MODEL_CLASSES.keys()),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite_output_dir", action="store_true", help="Overwrite the content of the output directory"
|
||||
)
|
||||
|
||||
parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
|
||||
parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")
|
||||
parser.add_argument(
|
||||
"--evaluate_during_training", action="store_true", help="Rul evaluation during training at each logging step."
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--per_device_train_batch_size", default=8, type=int, help="Batch size per GPU/CPU for training."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--per_device_eval_batch_size",
|
||||
default=8,
|
||||
type=int,
|
||||
help="Batch size per GPU/CPU for validation during training.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--per_device_test_batch_size",
|
||||
default=8,
|
||||
type=int,
|
||||
help="Batch size per GPU/CPU for evaluation after training.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
|
||||
)
|
||||
|
||||
parser.add_argument("--num_train_epochs", default=3, type=int, help="Total number of training epochs to perform.")
|
||||
parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
|
||||
parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
|
||||
parser.add_argument(
|
||||
"--max_seq_length",
|
||||
default=128,
|
||||
type=int,
|
||||
help="The maximum total input sequence length after tokenization. Sequences longer "
|
||||
"than this will be truncated, sequences shorter will be padded.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--doc_stride",
|
||||
default=128,
|
||||
type=int,
|
||||
help="When splitting up a long document into chunks, how much stride to take between chunks.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_query_length",
|
||||
default=64,
|
||||
type=int,
|
||||
help="The maximum number of tokens for the question. Questions longer than this will "
|
||||
"be truncated to this length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--do_lower_case", action="store_true", help="Set this flag if you are using an uncased model."
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tokenizer_name",
|
||||
default="",
|
||||
type=str,
|
||||
help="Pretrained tokenizer name or path if not the same as model_name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache_dir",
|
||||
default="",
|
||||
type=str,
|
||||
help="Where do you want to store the pre-trained models downloaded from s3",
|
||||
)
|
||||
parser.add_argument("--xla", action="store_true", help="Whether to use XLA (Accelerated Linear Algebra).")
|
||||
parser.add_argument("--amp", action="store_true", help="Whether to use AMP (Automatic Mixed Precision).")
|
||||
parser.add_argument(
|
||||
"--force_download",
|
||||
action="store_true",
|
||||
help="Whether to force download the weights from S3 (useful if the file is corrupted).",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--tpu",
|
||||
default=None,
|
||||
help="The Cloud TPU to use for training. This should be either the name "
|
||||
"used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 "
|
||||
"url.",
|
||||
)
|
||||
parser.add_argument("--num_tpu_cores", default="8", help="Total number of TPU cores to use.")
|
||||
parser.add_argument(
|
||||
"--gpus",
|
||||
default="0",
|
||||
help="Comma separated list of gpus devices. If only one, switch to single gpu strategy, if None takes all the gpus available.",
|
||||
)
|
||||
parser.add_argument("--no_cuda", action="store_true", help="Whether not to use CUDA when available")
|
||||
parser.add_argument(
|
||||
"--version_2_with_negative",
|
||||
action="store_true",
|
||||
help="If true, the SQuAD examples contain some that do not have an answer.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_steps",
|
||||
default=-1,
|
||||
help="If > 0: set total number of training steps to perform. Override num_train_epochs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gradient_accumulation_steps",
|
||||
default=1,
|
||||
help="Number of updates steps to accumulate before performing a backward/update pass.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup_steps", default=0, help="Linear warmup over warmup_steps.",
|
||||
)
|
||||
parser.add_argument("--max_grad_norm", default=1.0, help="Max gradient norm.")
|
||||
parser.add_argument("--logging_steps", default=500, type=int, help="Log every X updates.")
|
||||
parser.add_argument("--save_steps", default=500, type=int, help="Save checkpoint every X updates.")
|
||||
parser.add_argument(
|
||||
"--eval_all_checkpoints",
|
||||
action="store_true",
|
||||
help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n_best_size",
|
||||
default=20,
|
||||
type=int,
|
||||
help="The total number of n-best predictions to generate in the nbest_predictions.json output file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_answer_length",
|
||||
default=30,
|
||||
type=int,
|
||||
help="The maximum length of an answer that can be generated. This is needed because the start "
|
||||
"and end predictions are not conditioned on one another.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose_logging",
|
||||
action="store_true",
|
||||
help="If true, all of the warnings related to data processing will be printed. "
|
||||
"A number of warnings are expected for a normal SQuAD evaluation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--null_score_diff_threshold",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="If null_score - best_non_null is greater than the threshold predict null.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
)
|
||||
|
||||
if os.path.exists(args.output_dir) and args.do_train:
|
||||
if not args.overwrite_output_dir and bool(
|
||||
[file for file in os.listdir(args.output_dir) if "features" not in file]
|
||||
):
|
||||
raise ValueError(
|
||||
"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
|
||||
args.output_dir
|
||||
)
|
||||
)
|
||||
|
||||
if args.amp:
|
||||
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": True})
|
||||
|
||||
if args.tpu:
|
||||
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=args.tpu)
|
||||
tf.config.experimental_connect_to_cluster(resolver)
|
||||
tf.tpu.experimental.initialize_tpu_system(resolver)
|
||||
strategy = tf.distribute.experimental.TPUStrategy(resolver)
|
||||
args.n_device = args.num_tpu_cores
|
||||
elif len(args.gpus.split(",")) > 1:
|
||||
args.n_device = len([f"/gpu:{gpu}" for gpu in args.gpus.split(",")])
|
||||
strategy = tf.distribute.MirroredStrategy(devices=[f"/gpu:{gpu}" for gpu in args.gpus.split(",")])
|
||||
elif args.no_cuda:
|
||||
args.n_device = 1
|
||||
strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
|
||||
else:
|
||||
args.n_device = len(args.gpus.split(","))
|
||||
strategy = tf.distribute.OneDeviceStrategy(device="/gpu:" + args.gpus.split(",")[0])
|
||||
|
||||
logging.warning(
|
||||
"n_device: %s, distributed training: %s, 16-bits training: %s",
|
||||
args.n_device,
|
||||
bool(args.n_device > 1),
|
||||
args.amp,
|
||||
)
|
||||
|
||||
args.model_type = args.model_type.lower()
|
||||
config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
|
||||
|
||||
config = config_class.from_pretrained(
|
||||
args.config_name if args.config_name else args.model_name_or_path,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
)
|
||||
|
||||
logging.info("Training/evaluation parameters %s", args)
|
||||
|
||||
if args.do_train:
|
||||
tokenizer = tokenizer_class.from_pretrained(
|
||||
args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
|
||||
do_lower_case=args.do_lower_case,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
)
|
||||
|
||||
with strategy.scope():
|
||||
model = model_class.from_pretrained(
|
||||
args.model_name_or_path,
|
||||
from_pt=bool(".bin" in args.model_name_or_path),
|
||||
config=config,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
)
|
||||
model.layers[-1].activation = tf.keras.activations.softmax
|
||||
|
||||
train_batch_size = args.per_device_train_batch_size * args.n_device
|
||||
train_dataset, num_train_examples = load_and_cache_examples(args, tokenizer, evaluate=False)
|
||||
|
||||
train_dataset = train_dataset.batch(train_batch_size)
|
||||
train_dataset = train_dataset.prefetch(buffer_size=train_batch_size)
|
||||
train_dataset = strategy.experimental_distribute_dataset(train_dataset)
|
||||
|
||||
train(
|
||||
args, strategy, train_dataset, tokenizer, model, num_train_examples, train_batch_size,
|
||||
)
|
||||
|
||||
if not os.path.exists(args.output_dir):
|
||||
os.makedirs(args.output_dir)
|
||||
|
||||
logging.info("Saving model to %s", args.output_dir)
|
||||
|
||||
model.save_pretrained(args.output_dir)
|
||||
tokenizer.save_pretrained(args.output_dir)
|
||||
|
||||
if args.do_eval:
|
||||
tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
|
||||
checkpoints = []
|
||||
results = []
|
||||
|
||||
if args.eval_all_checkpoints:
|
||||
checkpoints = list(
|
||||
os.path.dirname(c)
|
||||
for c in sorted(
|
||||
glob.glob(args.output_dir + "/**/" + TF2_WEIGHTS_NAME, recursive=True),
|
||||
key=lambda f: int("".join(filter(str.isdigit, f)) or -1),
|
||||
)
|
||||
)
|
||||
|
||||
logging.info("Evaluate the following checkpoints: %s", checkpoints)
|
||||
|
||||
if len(checkpoints) == 0:
|
||||
checkpoints.append(args.output_dir)
|
||||
|
||||
for checkpoint in checkpoints:
|
||||
global_step = checkpoint.split("-")[-1] if re.match(".*checkpoint-[0-9]", checkpoint) else "final"
|
||||
|
||||
with strategy.scope():
|
||||
model = model_class.from_pretrained(checkpoint)
|
||||
results = evaluate(args, strategy, model, tokenizer, prefix=global_step)
|
||||
|
||||
result = dict((k + ("_{}".format(global_step) if global_step else ""), v) for k, v in results.items())
|
||||
results.update(result)
|
||||
|
||||
logger.info("Results: {}".format(results))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -555,10 +555,10 @@ def compute_predictions_logits(
|
||||
all_nbest_json[example.qas_id] = nbest_json
|
||||
|
||||
with open(output_prediction_file, "w") as writer:
|
||||
writer.write(json.dumps(all_predictions, indent=4) + "\n")
|
||||
writer.write(json.dumps(str(all_predictions), indent=4) + "\n")
|
||||
|
||||
with open(output_nbest_file, "w") as writer:
|
||||
writer.write(json.dumps(all_nbest_json, indent=4) + "\n")
|
||||
writer.write(json.dumps(str(all_nbest_json), indent=4) + "\n")
|
||||
|
||||
if version_2_with_negative:
|
||||
with open(output_null_log_odds_file, "w") as writer:
|
||||
|
||||
@@ -306,13 +306,15 @@ def squad_convert_examples_to_features(
|
||||
tqdm(
|
||||
p.imap(annotate_, examples, chunksize=32),
|
||||
total=len(examples),
|
||||
desc="convert squad examples to features",
|
||||
desc="Converting squad examples to features",
|
||||
)
|
||||
)
|
||||
|
||||
print("Converted {} examples into {} features".format(len(examples), len(features)))
|
||||
new_features = []
|
||||
unique_id = 1000000000
|
||||
example_index = 0
|
||||
for example_features in tqdm(features, total=len(features), desc="add example index and unique id"):
|
||||
for example_features in features:
|
||||
if not example_features:
|
||||
continue
|
||||
for example_feature in example_features:
|
||||
@@ -376,31 +378,34 @@ def squad_convert_examples_to_features(
|
||||
},
|
||||
)
|
||||
|
||||
return tf.data.Dataset.from_generator(
|
||||
gen,
|
||||
(
|
||||
{"input_ids": tf.int32, "attention_mask": tf.int32, "token_type_ids": tf.int32},
|
||||
{
|
||||
"start_position": tf.int64,
|
||||
"end_position": tf.int64,
|
||||
"cls_index": tf.int64,
|
||||
"p_mask": tf.int32,
|
||||
"is_impossible": tf.int32,
|
||||
},
|
||||
),
|
||||
(
|
||||
{
|
||||
"input_ids": tf.TensorShape([None]),
|
||||
"attention_mask": tf.TensorShape([None]),
|
||||
"token_type_ids": tf.TensorShape([None]),
|
||||
},
|
||||
{
|
||||
"start_position": tf.TensorShape([]),
|
||||
"end_position": tf.TensorShape([]),
|
||||
"cls_index": tf.TensorShape([]),
|
||||
"p_mask": tf.TensorShape([None]),
|
||||
"is_impossible": tf.TensorShape([]),
|
||||
},
|
||||
return (
|
||||
features,
|
||||
tf.data.Dataset.from_generator(
|
||||
gen,
|
||||
(
|
||||
{"input_ids": tf.int32, "attention_mask": tf.int32, "token_type_ids": tf.int32},
|
||||
{
|
||||
"start_position": tf.int64,
|
||||
"end_position": tf.int64,
|
||||
"cls_index": tf.int64,
|
||||
"p_mask": tf.int32,
|
||||
"is_impossible": tf.int32,
|
||||
},
|
||||
),
|
||||
(
|
||||
{
|
||||
"input_ids": tf.TensorShape([None]),
|
||||
"attention_mask": tf.TensorShape([None]),
|
||||
"token_type_ids": tf.TensorShape([None]),
|
||||
},
|
||||
{
|
||||
"start_position": tf.TensorShape([]),
|
||||
"end_position": tf.TensorShape([]),
|
||||
"cls_index": tf.TensorShape([]),
|
||||
"p_mask": tf.TensorShape([None]),
|
||||
"is_impossible": tf.TensorShape([]),
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -56,7 +56,16 @@ class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):
|
||||
}
|
||||
|
||||
|
||||
def create_optimizer(init_lr, num_train_steps, num_warmup_steps):
|
||||
def create_optimizer(
|
||||
init_lr,
|
||||
num_train_steps,
|
||||
num_warmup_steps,
|
||||
weight_decay=0.0,
|
||||
adam_epsilon=1e-6,
|
||||
beta_1=0.9,
|
||||
beta_2=0.999,
|
||||
exclude_from_weight_decay=("layer_norm", "bias"),
|
||||
):
|
||||
"""Creates an optimizer with learning rate schedule."""
|
||||
# Implements linear decay of the learning rate.
|
||||
learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(
|
||||
@@ -68,11 +77,11 @@ def create_optimizer(init_lr, num_train_steps, num_warmup_steps):
|
||||
)
|
||||
optimizer = AdamWeightDecay(
|
||||
learning_rate=learning_rate_fn,
|
||||
weight_decay_rate=0.01,
|
||||
beta_1=0.9,
|
||||
beta_2=0.999,
|
||||
epsilon=1e-6,
|
||||
exclude_from_weight_decay=["layer_norm", "bias"],
|
||||
weight_decay_rate=weight_decay,
|
||||
beta_1=beta_1,
|
||||
beta_2=beta_2,
|
||||
epsilon=adam_epsilon,
|
||||
exclude_from_weight_decay=exclude_from_weight_decay,
|
||||
)
|
||||
return optimizer
|
||||
|
||||
|
||||
@@ -998,7 +998,7 @@ class PreTrainedTokenizer(object):
|
||||
for key, value in batch_outputs.items():
|
||||
|
||||
padded_value = value
|
||||
if key != "input_len":
|
||||
if key != "input_len" and self._pad_token is not None:
|
||||
# Padding handle
|
||||
padded_value = [
|
||||
v + [self.pad_token_id if key == "input_ids" else 1] * (max_seq_len - len(v))
|
||||
|
||||
Reference in New Issue
Block a user