import os
import sys
import re
import collections
import numpy as np
import scipy
import json
import itertools
import pickle
import gc
import gzip
import argparse
def tokenize(sent):
'''Return the tokens of a sentence including punctuation.
>>> tokenize('Bob dropped the apple. Where is the apple?')
['Bob', 'dropped', 'the', 'apple', '.', 'Where', 'is', 'the', 'apple', '?']
'''
return re.findall('(?:\w+)|\S',sent)
def list_to_map(l):
'''Convert a list of values to a map from values to indices'''
return {val:i for i,val in enumerate(l)}
def parse_stories(lines):
'''
Parse stories provided in the bAbi tasks format, with knowledge graph.
'''
data = []
story = []
for line in lines:
if line[-1] == "\n":
line = line[:-1]
nid, line = line.split(' ', 1)
nid = int(nid)
if nid == 1:
story = []
questions = []
if '\t' in line:
q, apre = line.split('\t')[:2]
a = apre.split(',')
q = tokenize(q)
substory = [x for x in story if x]
data.append((substory, q, a))
story.append('')
else:
line, graph = line.split('=', 1)
sent = tokenize(line)
graph_parsed = json.loads(graph)
story.append((sent, graph_parsed))
return data
def get_stories(taskname):
with open(taskname, 'r') as f:
lines = f.readlines()
return parse_stories(lines)
def get_max_sentence_length(stories):
return max((max((len(sentence) for (sentence, graph) in sents_graphs)) for (sents_graphs, query, answer) in stories))
def get_max_query_length(stories):
return max((len(query) for (sents_graphs, query, answer) in stories))
def get_max_num_queries(stories):
return max((len(queries) for (sents_graphs, query, answer) in stories))
def get_max_nodes_per_iter(stories):
result = 0
for (sents_graphs, query, answer) in stories:
prev_nodes = set()
for (sentence, graph) in sents_graphs:
cur_nodes = set(graph["nodes"])
new_nodes = len(cur_nodes - prev_nodes)
if new_nodes > result:
result = new_nodes
prev_nodes = cur_nodes
return result
def get_buckets(stories, max_ignore_unbatched=100, max_pad_amount=25):
sentencecounts = [len(sents_graphs) for (sents_graphs, query, answer) in stories]
countpairs = sorted(collections.Counter(sentencecounts).items())
buckets = []
smallest_left_val = 0
num_unbatched = max_ignore_unbatched
for val,ct in countpairs:
num_unbatched += ct
if val - smallest_left_val > max_pad_amount or num_unbatched > max_ignore_unbatched:
buckets.append(val)
smallest_left_val = val
num_unbatched = 0
if buckets[-1] != countpairs[-1][0]:
buckets.append(countpairs[-1][0])
return buckets
PAD_WORD = "<PAD>"
def get_wordlist(stories):
words = [PAD_WORD] + sorted(list(set((word
for (sents_graphs, query, answer) in stories
for wordbag in itertools.chain((s for s,g in sents_graphs), [query])
for word in wordbag ))))
wordmap = list_to_map(words)
return words, wordmap
def get_answer_list(stories):
words = sorted(list(set(word for (sents_graphs, query, answer) in stories for word in answer)))
wordmap = list_to_map(words)
return words, wordmap
def pad_story(story, num_sentences, sentence_length):
def pad(lst,dlen,pad):
return lst + [pad]*(dlen - len(lst))
sents_graphs, query, answer = story
padded_sents_graphs = [(pad(s,sentence_length,PAD_WORD), g) for s,g in sents_graphs]
padded_query = pad(query,sentence_length,PAD_WORD)
sentgraph_padding = (pad([],sentence_length,PAD_WORD), padded_sents_graphs[-1][1])
return (pad(padded_sents_graphs, num_sentences, sentgraph_padding), padded_query, answer)
def get_unqualified_id(s):
return s.split("#")[0]
def get_graph_lists(stories):
node_words = sorted(list(set(get_unqualified_id(node)
for (sents_graphs, query, answer) in stories
for sent,graph in sents_graphs
for node in graph["nodes"])))
nodemap = list_to_map(node_words)
edge_words = sorted(list(set(get_unqualified_id(edge["type"])
for (sents_graphs, query, answer) in stories
for sent,graph in sents_graphs
for edge in graph["edges"])))
edgemap = list_to_map(edge_words)
return node_words, nodemap, edge_words, edgemap
def convert_graph(graphs, nodemap, edgemap, new_nodes_per_iter, dynamic=True):
num_node_ids = len(nodemap)
num_edge_types = len(edgemap)
full_size = len(graphs)*new_nodes_per_iter + 1
prev_size = 1
processed_nodes = []
index_map = {}
all_num_nodes = []
all_node_ids = []
all_node_strengths = []
all_edges = []
if not dynamic:
processed_nodes = list(nodemap.keys())
index_map = nodemap.copy()
prev_size = num_node_ids
full_size = prev_size
new_nodes_per_iter = 0
for g in graphs:
active_nodes = g["nodes"]
active_edges = g["edges"]
new_nodes = [e for e in active_nodes if e not in processed_nodes]
num_new_nodes = len(new_nodes)
if not dynamic:
assert num_new_nodes == 0, "Cannot create more nodes in non-dynamic mode!\n{}".format(graphs)
new_node_strengths = np.zeros([new_nodes_per_iter], np.float32)
new_node_strengths[:num_new_nodes] = 1.0
new_node_ids = np.zeros([new_nodes_per_iter, num_node_ids], np.float32)
for i, node in enumerate(new_nodes):
new_node_ids[i,nodemap[get_unqualified_id(node)]] = 1.0
index_map[node] = prev_size + i
next_edges = np.zeros([full_size, full_size, num_edge_types])
for edge in active_edges:
next_edges[index_map[edge["from"]],
index_map[edge["to"]],
edgemap[get_unqualified_id(edge["type"])]] = 1.0
processed_nodes.extend(new_nodes)
prev_size += new_nodes_per_iter
all_num_nodes.append(num_new_nodes)
all_node_ids.append(new_node_ids)
all_edges.append(next_edges)
all_node_strengths.append(new_node_strengths)
return np.stack(all_num_nodes), np.stack(all_node_strengths), np.stack(all_node_ids), np.stack(all_edges)
def convert_story(story, wordmap, answer_map, graph_node_map, graph_edge_map, new_nodes_per_iter, dynamic=True):
"""
Converts a story in format
([(sentence, graph)], [(index, question_arr, answer)])
to a consolidated story in format
(sentence_arr, [graph_arr_dict], [(index, question_arr, answer)])
and also replaces words according to the input maps
"""
sents_graphs, query, answer = story
sentence_arr = [[wordmap[w] for w in s] for s,g in sents_graphs]
graphs = convert_graph([g for s,g in sents_graphs], graph_node_map, graph_edge_map, new_nodes_per_iter, dynamic)
query_arr = [wordmap[w] for w in query]
answer_arr = [answer_map[w] for w in answer]
return (sentence_arr, graphs, query_arr, answer_arr)
def process_story(s,bucket_len):
return convert_story(pad_story(s, bucket_len, sentence_length), wordmap, answer_map, graph_node_map, graph_edge_map, new_nodes_per_iter, dynamic)
def bucket_stories(stories, buckets, wordmap, answer_map, graph_node_map, graph_edge_map, sentence_length, new_nodes_per_iter, dynamic=True):
return [ [process_story(story,bmax) for story in stories if bstart < len(story[0]) <= bmax]
for bstart, bmax in zip([0]+buckets,buckets)]
def prepare_stories(stories, dynamic=True):
sentence_length = max(get_max_sentence_length(stories), get_max_query_length(stories))
buckets = get_buckets(stories)
wordlist, wordmap = get_wordlist(stories)
anslist, ansmap = get_answer_list(stories)
new_nodes_per_iter = get_max_nodes_per_iter(stories)
graph_node_list, graph_node_map, graph_edge_list, graph_edge_map = get_graph_lists(stories)
bucketed = bucket_stories(stories, buckets, wordmap, ansmap, graph_node_map, graph_edge_map, sentence_length, new_nodes_per_iter, dynamic)
return sentence_length, new_nodes_per_iter, buckets, wordlist, anslist, graph_node_list, graph_edge_list, bucketed
def print_batch(story, wordlist, anslist, file=sys.stdout):
sents, query, answer = story
for batch,(s,q,a) in enumerate(zip(sents,query,answer)):
file.write("Story {}\n".format(batch))
for sent in s:
file.write(" ".join([wordlist[word] for word in sent]) + "\n")
file.write(" ".join(wordlist[word] for word in q) + "\n")
file.write(" ".join(anslist[word] for word in a.nonzero()[1]) + "\n")
MetadataList = collections.namedtuple("MetadataList", ["sentence_length", "new_nodes_per_iter", "buckets", "wordlist", "anslist", "graph_node_list", "graph_edge_list"])
PreppedStory = collections.namedtuple("PreppedStory", ["converted", "sentences", "query", "answer"])
def generate_metadata(stories, dynamic=True):
sentence_length = max(get_max_sentence_length(stories), get_max_query_length(stories))
buckets = get_buckets(stories)
wordlist, wordmap = get_wordlist(stories)
anslist, ansmap = get_answer_list(stories)
new_nodes_per_iter = get_max_nodes_per_iter(stories)
graph_node_list, graph_node_map, graph_edge_list, graph_edge_map = get_graph_lists(stories)
metadata = MetadataList(sentence_length, new_nodes_per_iter, buckets, wordlist, anslist, graph_node_list, graph_edge_list)
return metadata
def preprocess_stories(stories, savedir, dynamic=True, metadata_file=None):
if metadata_file is None:
metadata = generate_metadata(stories, dynamic)
else:
with open(metadata_file,'rb') as f:
metadata = pickle.load(f)
buckets = get_buckets(stories)
sentence_length, new_nodes_per_iter, old_buckets, wordlist, anslist, graph_node_list, graph_edge_list = metadata
metadata = metadata._replace(buckets=buckets)
if not os.path.exists(savedir):
os.makedirs(savedir)
with open(os.path.join(savedir,'metadata.p'),'wb') as f:
pickle.dump(metadata, f)
bucketed_files = [[] for _ in buckets]
for i,story in enumerate(stories):
bucket_idx, cur_bucket = next(((i,bmax) for (i,(bstart, bmax)) in enumerate(zip([0]+buckets,buckets))
if bstart < len(story[0]) <= bmax), (None,None))
assert cur_bucket is not None, "Couldn't put story of length {} into buckets {}".format(len(story[0]), buckets)
bucket_dir = os.path.join(savedir, "bucket_{}".format(cur_bucket))
if not os.path.exists(bucket_dir):
os.makedirs(bucket_dir)
story_fn = os.path.join(bucket_dir, "story_{}.pz".format(i))
sents_graphs, query, answer = story
sents = [s for s,g in sents_graphs]
cvtd = convert_story(pad_story(story, cur_bucket, sentence_length), list_to_map(wordlist), list_to_map(anslist), list_to_map(graph_node_list), list_to_map(graph_edge_list), new_nodes_per_iter, dynamic)
prepped = PreppedStory(cvtd, sents, query, answer)
with gzip.open(story_fn, 'wb') as zf:
pickle.dump(prepped, zf)
bucketed_files[bucket_idx].append(os.path.relpath(story_fn, savedir))
gc.collect() # we don't want to use too much memory, so try to clean it up
with open(os.path.join(savedir,'file_list.p'),'wb') as f:
pickle.dump(bucketed_files, f)
def main(file, dynamic, metadata_file=None):
stories = get_stories(file)
dirname, ext = os.path.splitext(file)
preprocess_stories(stories, dirname, dynamic, metadata_file)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Parse a graph file')
parser.add_argument("file", help="Graph file to parse")
parser.add_argument("--static", dest="dynamic", action="store_false", help="Don't use dynamic nodes")
parser.add_argument("--metadata-file", default=None, help="Use this particular metadata file instead of building it from scratch")
args = vars(parser.parse_args())
main(**args)
