Spaces:
Runtime error
Runtime error
File size: 3,738 Bytes
7900c16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
"""
This script provides an example to wrap TencentPretrain for Chinese machine reading comprehension inference.
"""
import sys
import os
import argparse
import torch
import torch.nn as nn
tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(tencentpretrain_dir)
from tencentpretrain.utils.config import load_hyperparam
from tencentpretrain.utils.constants import *
from tencentpretrain.utils.tokenizers import *
from tencentpretrain.model_loader import load_model
from tencentpretrain.opts import infer_opts
from finetune.run_cmrc import *
def main():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
infer_opts(parser)
parser.add_argument("--vocab_path", default=None, type=str,
help="Path of the vocabulary file.")
parser.add_argument("--spm_model_path", default=None, type=str,
help="Path of the sentence piece model.")
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.")
args = parser.parse_args()
# Load the hyperparameters from the config file.
args = load_hyperparam(args)
# Build tokenizer.
args.tokenizer = CharTokenizer(args)
# Build model and load parameters.
model = MachineReadingComprehension(args)
model = load_model(model, args.load_model_path)
# For simplicity, we use DataParallel wrapper to use multiple GPUs.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
if torch.cuda.device_count() > 1:
print("{} GPUs are available. Let's use them.".format(torch.cuda.device_count()))
model = torch.nn.DataParallel(model)
dataset, examples = read_dataset(args, args.test_path)
src = torch.LongTensor([sample[0] for sample in dataset])
seg = torch.LongTensor([sample[1] for sample in dataset])
start_position = torch.LongTensor([sample[2] for sample in dataset])
end_position = torch.LongTensor([sample[3] for sample in dataset])
batch_size = args.batch_size
instances_num = len(dataset)
print("The number of prediction instances: ", instances_num)
model.eval()
with open(args.prediction_path, mode="w", encoding="utf-8") as f:
start_prob_all, end_prob_all = [], []
for i, (src_batch, seg_batch, start_position_batch, end_position_batch) in enumerate(batch_loader(batch_size, src, seg, start_position, end_position)):
src_batch = src_batch.to(device)
seg_batch = seg_batch.to(device)
start_position_batch = start_position_batch.to(device)
end_position_batch = end_position_batch.to(device)
with torch.no_grad():
loss, start_logits, end_logits = model(src_batch, seg_batch, start_position_batch, end_position_batch)
start_prob = nn.Softmax(dim=1)(start_logits)
end_prob = nn.Softmax(dim=1)(end_logits)
for j in range(start_prob.size()[0]):
start_prob_all.append(start_prob[j])
end_prob_all.append(end_prob[j])
pred_answers = get_answers(dataset, start_prob_all, end_prob_all)
output = {}
for i in range(len(examples)):
question_id = examples[i][2]
start_pred_pos = pred_answers[i][1]
end_pred_pos = pred_answers[i][2]
prediction = examples[i][0][start_pred_pos: end_pred_pos + 1]
output[question_id] = prediction
f.write(json.dumps(output, indent=4, ensure_ascii=False) + "\n")
if __name__ == "__main__":
main()
|