def get_nsp_score_batch(nsp_predictor, predictions):
"""
Get NSP scores of a batch.
"""
import argparse
from collections import namedtuple
from readers.nsp_reader import NSPReader
from utils.args import parse_args
from tasks.next_sentence_prediction import NextSentencePrediction
parser = argparse.ArgumentParser()
NextSentencePrediction.add_cmdline_args(parser)
parser.add_argument("--num_samples", type=int, default=None)
parser.add_argument("--config_path", type=str, required=True)
parser.add_argument("--mem_efficient", type=str2bool, default=False)
args = parse_args(parser, allow_unknown=True)
args.load(args.config_path)
if not args.mem_efficient:
if args.num_samples:
args.batch_size *= args.num_samples
if args.latent_type_size:
args.batch_size *= args.latent_type_size
args.tokenized_input = True
reader = NSPReader(args)
def __reader__():
headers = ["src", "tgt", "data_id"]
Example = namedtuple("Example", headers)
for i, info in enumerate(predictions):
context = post_process_context(info["context_token_ids"], reader, merge=False)
context_tokenized_input = " [SEP] ".join(" ".join(utt) for utt in context)
_, response = post_process_response(info["response_token_ids"], reader, merge=False)
response_tokenized_input = " ".join(response)
example = Example(src=context_tokenized_input, tgt=response_tokenized_input, data_id=i)
record = reader._convert_example_to_record(example, is_infer=True)
yield record
return
generator = reader.data_generator(
reader=__reader__,
is_infer=True,
phase="test",
)
steps = 0
for data in generator():
outputs = nsp_predictor(data)
for probs, data_id in zip(outputs[0], outputs[-1]):
data_id = data_id[0]
info = predictions[data_id]
info["nsp_score"] = float(probs[1])
return
def setup_args():
"""Setup arguments."""
parser = argparse.ArgumentParser()
group = parser.add_argument_group("Model")
group.add_argument("--init_from_ckpt", type=str, default="")
group.add_argument("--vocab_size", type=int, default=8001)
group.add_argument("--latent_type_size", type=int, default=20)
group.add_argument("--num_layers", type=int, default=24)
group = parser.add_argument_group("Task")
group.add_argument("--is_cn", type=str2bool, default=False)
args, _ = parser.parse_known_args()
NSPReader.add_cmdline_args(parser)
args = parse_args(parser)
args.batch_size *= args.latent_type_size
#print(json.dumps(args, indent=2))
return args
def interact(args):
"""Inference main function."""
plato_reader = PlatoReader(args)
nsp_reader = NSPReader(args)
if args.num_layers == 24:
n_head = 16
hidden_size = 1024
elif args.num_layers == 32:
n_head = 32
hidden_size = 2048
else:
raise ValueError('The pre-trained model only support 24 or 32 layers, '
'but received num_layers=%d.' % args.num_layers)
model = Plato2InferModel(nsp_reader, args.num_layers, n_head, hidden_size)
load_params(model, args.init_from_ckpt)
model.eval()
Example = namedtuple("Example", ["src", "data_id"])
context = []
start_info = "Enter [EXIT] to quit the interaction, [NEXT] to start a new conversation."
cprint(start_info, "yellow", attrs=["bold"])
while True:
user_utt = input(colored("[Human]: ", "red", attrs=["bold"])).strip()
if user_utt == "[EXIT]":
break
elif user_utt == "[NEXT]":
context = []
cprint(start_info, "yellow", attrs=["bold"])
else:
context.append(user_utt)
example = Example(src=" [SEP] ".join(context), data_id=0)
record = plato_reader._convert_example_to_record(
example, is_infer=True)
data = plato_reader._pad_batch_records([record], is_infer=True)
inputs = gen_inputs(data, args.latent_type_size)
inputs['tgt_ids'] = inputs['tgt_ids'].astype('int64')
pred = model(inputs)[0]
bot_response = pred["response"]
print(
colored(
"[Bot]:", "blue", attrs=["bold"]),
colored(
bot_response, attrs=["bold"]))
context.append(bot_response)
return
def __init__(self, args):
super(NextSentencePrediction, self).__init__(args)
self.reader = NSPReader(args)
return
def add_cmdline_args(cls, parser):
group = NSPReader.add_cmdline_args(parser)
return group
def add_cmdline_args(cls, parser):
"""Add cmdline argurments."""
group = NSPReader.add_cmdline_args(parser)
return group