yuyan-dialogue / inference.py

席亚东

fix the bug in inference.py

16bf127 over 1 year ago

No virus

7.56 kB

	#!/usr/bin/env python3 -u

	from collections import namedtuple

	import math
	import torch
	from torch.nn.utils.rnn import pad_sequence

	from fairseq import checkpoint_utils, options, tasks, utils
	from eet.fairseq.transformer import EETTransformerDecoder

	Batch = namedtuple('Batch', 'ids src_tokens src_lengths')

	def make_batches(lines, task, max_positions, encode_fn):

	tokens = [task.source_dictionary.encode_line(encode_fn(line),
	add_if_not_exist=False,
	append_eos=False,
	reverse_order=True).long()
	for line in lines]
	lengths = [t.numel() for t in tokens]
	tokens = pad_sequence(tokens, batch_first=True,
	padding_value=1).flip(dims=(1,))

	return Batch(ids=torch.arange(len(tokens)),
	src_tokens=tokens,
	src_lengths=torch.tensor(lengths))

	def encode_fn(x_str):
	x_str = "</s> " + x_str
	return x_str


	def decode_fn(x):
	x = x.replace(" ", "")
	return x

	def eos_token_filter(sent):
	return True


	def post_precess(line):

	if "<" in line:
	line = line.split("<")[0]
	return line


	class Inference(object):

	def __init__(self, model_path, data_path, eet_batch_size):

	parser = options.get_generation_parser(default_task="language_modeling")
	args = options.parse_args_and_arch(parser)
	args.data = data_path
	args.path = model_path
	self.args = args

	# generate parameter
	args.beam = 1 # don't change
	args.min_len = 5
	args.max_len_b = 30
	args.lenpen = 1.0
	args.sampling = True
	# args.sampling_topp = 0.7
	args.sampling_topk = 10
	args.temperature = 0.8
	args.no_repeat_ngram_size = 1
	args.fp16 = True

	# Setup task, e.g., translation
	task = tasks.setup_task(args)
	self.task = task
	# Set dictionaries
	self.src_dict = task.source_dictionary
	self.tgt_dict = task.target_dictionary

	use_cuda = torch.cuda.is_available() and not args.cpu
	self.use_cuda = use_cuda

	model_path = args.path
	checkpoint = torch.load(model_path.replace("best.pt", "best_part_1.pt"))
	checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_2.pt")))
	checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_3.pt")))
	torch.save(checkpoint, model_path)

	state = torch.load(args.path, map_location=torch.device("cpu"))
	cfg_args = eval(str(state["cfg"]))["model"]
	del cfg_args["_name"]
	keys_list = []
	values_list = []
	for key,value in cfg_args.items() :
	keys_list.append(key)
	values_list.append(value)
	Model_args = namedtuple("Model_args", keys_list)
	model_args = Model_args._make(values_list)
	del state

	eet_seq_len = 512 # max seqence length
	eet_batch_size = eet_batch_size
	data_type = torch.float16
	eet_config = {"data_type":data_type,
	"max_batch":eet_batch_size,
	"full_seq_len":eet_seq_len}
	print(model_args)
	eet_model = EETTransformerDecoder.from_fairseq_pretrained(model_id_or_path = args.path,
	dictionary = self.src_dict,args=model_args,
	config = eet_config,
	no_encoder_attn = True)
	self.models = [eet_model]
	# Initialize generator
	self.generator = task.build_generator(self.models, args)

	# Load alignment dictionary for unknown word replacement
	# (None if no unknown word replacement, empty if no path to align dictionary)
	self.align_dict = utils.load_align_dict(args.replace_unk)

	self.max_positions = 1024
	self.eos_index = self.tgt_dict.eos()
	self.pad_index = self.tgt_dict.pad()

	def __call__(self, inputs, append_right_eos=True):

	results = []
	start_id = 0

	batch = make_batches(inputs, self.task, self.max_positions, encode_fn)
	inputs_str = inputs

	src_tokens = batch.src_tokens
	src_lengths = batch.src_lengths
	# a new paragraph always
	if src_tokens[0][-1].item() != self.eos_index and append_right_eos:
	src_tokens = torch.cat([src_tokens, src_tokens.new_ones(src_tokens.size(0), 1) * self.eos_index], dim=1)
	src_lengths += 1
	if self.use_cuda:
	src_tokens = src_tokens.cuda()
	src_lengths = src_lengths.cuda()
	sample = {
	'net_input': {
	'src_tokens': src_tokens,
	'src_lengths': src_lengths,
	},
	}

	translations = self.task.inference_step(self.generator, self.models, sample)

	for i, (id, hypos) in enumerate(zip(batch.ids.tolist(), translations)):
	results.append((start_id + id, src_tokens[i], hypos))

	# sort output to match input order
	final_results = []
	for id, src_tokens, hypos in sorted(results, key=lambda x: x[0]):
	# Process top predictions
	tmp_res = []
	for hypo in hypos[:min(len(hypos), self.args.nbest)]:
	hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
	hypo_tokens=hypo['tokens'].int().cpu()[len(src_tokens)-1:],
	src_str=None,
	alignment=hypo['alignment'],
	align_dict=self.align_dict,
	tgt_dict=self.tgt_dict)

	detok_hypo_str = decode_fn(hypo_str)
	if eos_token_filter(detok_hypo_str):
	detok_hypo_str = post_precess(detok_hypo_str)
	score = hypo['score'] / math.log(2) # convert to base 2
	tmp_res.append([detok_hypo_str, score])
	final_results.append(tmp_res)
	return final_results




	class Dialogue(object):
	def __init__(self, inference_model=None, max_dialogue_history=6):

	self.inference_model = inference_model
	self.max_dialogue_history = max_dialogue_history
	self.dialogue_history = []

	def get_repsonse(self, input_text):
	self.dialogue_history.append(input_text.strip())
	model_inp = ""
	for idx, x in enumerate(self.dialogue_history[-self.max_dialogue_history:]):
	if idx % 2 == 0:
	model_inp += " <0> " + " ".join(list(x))
	else:
	model_inp += " <1> " + " ".join(list(x))
	if idx % 2 == 0:
	model_inp += " <1>"
	else:
	model_inp += " <0>"
	# generate 5 candidates
	text = self.inference_model([model_inp]*5, append_right_eos=False)
	response = [x[0][0] for x in text]
	# response rank according to length
	response = sorted(response, key=lambda x:len(set(x)))
	# overlap-score
	overlap = [[len(set(x) & set(model_inp)) * len(x), x] for x in response[-4:-1]]
	overlap = sorted(overlap, key=lambda x:x[0])
	final_response = overlap[-2][1]
	self.dialogue_history.append(final_response)
	return final_response

	def clear_dialogue_history(self):
	self.dialogue_history = []