def collect_babi_slot_values(in_babi_root):
dataset_files = get_files_list(in_babi_root, 'task1-API-calls')
babi_files = [(filename, read_task(filename))
for filename in dataset_files]
full_babi = reduce(lambda x, y: x + y[1], babi_files, [])
slots_map = extract_slot_values(full_babi)
return reduce(lambda x, y: list(x) + list(y), slots_map.values(), [])
def plus_single_task(in_task, slot_values):
slots_map = extract_slot_values(in_task) \
if slot_values is None \
else slot_values
babi_plus = map(
lambda dialogue: augment_dialogue(dialogue, slots_map.values()),
in_task)
return babi_plus
def plus_dataset(in_src_root, in_result_size):
dataset_files = get_files_list(in_src_root, 'task1-API-calls')
babi_files = [(filename, read_task(filename))
for filename in dataset_files]
full_babi = reduce(lambda x, y: x + y[1], babi_files, [])
slots_map = extract_slot_values(full_babi)
babi_plus = defaultdict(lambda: [])
result_size = in_result_size if in_result_size else len(babi_files)
for task_name, task in babi_files:
for dialogue_index, dialogue in zip(xrange(result_size), cycle(task)):
babi_plus[task_name].append(
augment_dialogue(dialogue, slots_map.values()))
return babi_plus
def main(in_config, in_babi_file, in_result_file):
init(in_config)
task = read_task(in_babi_file)
slot_values = extract_slot_values(task)
babi_plus_dialogues = plus_single_task(task, slot_values)
utterances, tags, pos = [], [], []
for dialogue in babi_plus_dialogues:
for turn in dialogue:
if turn['agent'] == 'user':
utterances.append(turn['text'].split())
tags.append(turn['tags'])
pos.append(turn['pos'])
result = pd.DataFrame({'utterance': utterances, 'tags': tags, 'pos': pos})
result.to_json(in_result_file)
print_stats()
def configure_argument_parser():
parser = ArgumentParser(description='generate bAbI+ data')
parser.add_argument('babi_file', help='file with bAbI Dialogs')
parser.add_argument('babi_plus_root', help='output folder')
parser.add_argument('--output_format',
default='babi',
help='format of output dialogues [babi/babble]')
parser.add_argument(
'--result_size',
type=int,
default=None,
help='size of generated dataset [default=input dataset size]')
parser.add_argument('--config',
default=DEFAULT_CONFIG_FILE,
help='dicustom disfluency config (json file)')
return parser
if __name__ == '__main__':
parser = configure_argument_parser()
args = parser.parse_args()
init(args.config)
task = read_task(args.babi_file)
slot_values = extract_slot_values(task)
task_name = path.basename(args.babi_file)
babi_plus_dialogues = plus_single_task(task, slot_values)
save_function = locals()['save_' + args.output_format]
save_function({task_name: babi_plus_dialogues}, args.babi_plus_root)
print_stats()
