data_mode = ["dev"]
for dm in data_mode:
dpath = os.path.join(args.input_path, data_format.format(dm))
data = dev_reader._read(dpath) if dm == "dev" else train_reader._read(
dpath)
print("Save data to {}.".format(
os.path.join(args.output_dir,
"tmspan_cached_roberta_{}.pkl".format(dm))))
with open(
os.path.join(args.output_dir,
"tmspan_cached_roberta_{}.pkl".format(dm)),
"wb") as f:
pickle.dump(data, f)
else:
dev_reader = DropReader(tokenizer, args.passage_length_limit,
args.question_length_limit)
train_reader = DropReader(tokenizer,
args.passage_length_limit,
args.question_length_limit,
skip_when_all_empty=[
"passage_span",
"question_span",
"addition_subtraction",
"counting",
])
data_format = "drop_dataset_{}.json"
data_mode = ["train"]
for dm in data_mode:
dpath = os.path.join(args.input_path, data_format.format(dm))
print("Load from pre path {}.".format(args.pre_path))
network.load_state_dict(torch.load(args.pre_path))
print("Load data from {}.".format(args.inf_path))
tokenizer = RobertaTokenizer.from_pretrained(args.roberta_model)
if args.tag_mspan:
inf_iter = TDropBatchGen(
args, tokenizer,
TDropReader(tokenizer,
passage_length_limit=463,
question_length_limit=46)._read(args.inf_path))
else:
inf_iter = DropBatchGen(
args, tokenizer,
DropReader(tokenizer,
passage_length_limit=463,
question_length_limit=46)._read(args.inf_path))
print("Start inference...")
result = {}
network.eval()
with torch.no_grad():
for batch in tqdm(inf_iter):
output_dict = network(**batch)
for i in range(len(output_dict["question_id"])):
if "arithmetic_expression" in output_dict["answer"][i].keys():
result[output_dict["question_id"][i]] = {
"predicted_answer":
output_dict["answer"][i]["predicted_answer"],
"answer_type":
output_dict["answer"][i]["answer_type"],
print("Build bert model.")
bert_model = RobertaModel(RobertaConfig().from_pretrained(args.roberta_model))
print("Build Drop model.")
network = NumericallyAugmentedBertNet(bert_model,
hidden_size=bert_model.config.hidden_size,
dropout_prob=0.0,
use_gcn=args.use_gcn,
gcn_steps=args.gcn_steps)
if args.cuda: network.cuda()
print("Load from pre path {}.".format(args.pre_path))
network.load_state_dict(torch.load(args.pre_path))
print("Load data from {}.".format(args.inf_path))
tokenizer = RobertaTokenizer.from_pretrained(args.roberta_model)
inf_iter = DropBatchGen(args, tokenizer, DropReader(tokenizer, passage_length_limit=463, question_length_limit=46)._read(args.inf_path))
print("Start inference...")
result = {}
network.eval()
with torch.no_grad():
for batch in tqdm(inf_iter):
output_dict = network(**batch)
for i in range(len(output_dict["question_id"])):
result[output_dict["question_id"][i]] = output_dict["answer"][i]["predicted_answer"]
with open(args.dump_path, "w", encoding="utf8") as f:
json.dump(result, f)
import argparse
from pytorch_transformers.tokenization_roberta import RobertaTokenizer
from mspan_roberta_gcn.drop_roberta_dataset import DropReader
parser = argparse.ArgumentParser()
parser.add_argument("--input_path", type=str)
parser.add_argument("--output_dir", type=str)
parser.add_argument("--passage_length_limit", type=int, default=463)
parser.add_argument("--question_length_limit", type=int, default=46)
args = parser.parse_args()
tokenizer = RobertaTokenizer.from_pretrained(args.input_path +
"/roberta.large")
dev_reader = DropReader(tokenizer, args.passage_length_limit,
args.question_length_limit)
train_reader = DropReader(tokenizer,
args.passage_length_limit,
args.question_length_limit,
skip_when_all_empty=[
"passage_span",
"question_span",
"addition_subtraction",
"counting",
])
data_format = "drop_dataset_{}.json"
data_mode = ["train"]
for dm in data_mode: