Compare commits

...
Author SHA1 Message Date
Lysandre ffc522e1a1 Remaining arguments 2020-01-29 10:14:12 -05:00
Lysandre 459c7707f0 Data loading & Training & Evaluation 2020-01-28 16:06:05 -05:00
Lysandre 97debe4723 run tf squad training 2020-01-28 16:05:39 -05:00
Lysandre 9a9e68e7b5 Data loading 2020-01-28 16:05:39 -05:00
Lysandre 74577316b5 Revamping the GLUE script + init SQuAD script 2020-01-28 16:05:39 -05:00
Lysandre 2b6f10e4e3 Force download option 2020-01-28 16:05:39 -05:00
Lysandre 59988ba723 Better script for GLUE in TensorFlow 2020-01-28 16:05:39 -05:00
6 changed files with 1060 additions and 123 deletions
+311 -87
View File
@@ -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()
+699
View File
@@ -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:
+32 -27
View File
@@ -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([]),
},
),
),
)
+15 -6
View File
@@ -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
+1 -1
View File
@@ -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))