"""

A QAModel embeds ClozeQuery (see query.py) and their respective support to answer the query by selecting one of the

specified answer candidates. This is an implementation of the system described in [to appear].

"""

def __init__(self, size, batch_size, vocab_size, answer_vocab_size, max_length, is_train=True, learning_rate=1e-2,

composition="GRU", max_hops=0, devices=None, keep_prob=1.0):

"""

:param size: size of hidden states

:param batch_size: initial batch_size (adapts automatically)

:param vocab_size: size of input vocabulary (vocabulary of contexts)

:param answer_vocab_size: size of answer (candidates) vocabulary

:param max_length: maximum length of an individual context

:param is_train:

:param learning_rate:

:param composition: "GRU", "LSTM", "BiGRU" are possible

:param max_hops: maximum number of hops, can be set manually to something lower by assigning a different value

to variable (self.)num_hops which is initialized with max_hops

:param devices: defaults to ["/cpu:0"], but can be a list of up to 3 devices. The model is automatically

partitioned into the different devices.

:param keep_prob: 1.0-dropout rate, that is applied to the input embeddings

"""

self._vocab_size = vocab_size

self._max_length = max_length

self._size = size

self._batch_size = batch_size

self._is_train = is_train

self._composition = composition

self._max_hops = max_hops

self._device0 = devices[0] if devices is not None else "/cpu:0"

self._device1 = devices[1 % len(devices)] if devices is not None else "/cpu:0"

self._device2 = devices[2 % len(devices)] if devices is not None else "/cpu:0"

self._init = tf.random_normal_initializer(0.0, 0.1)

with tf.device(self._device0):

with tf.variable_scope(self.name(), initializer=tf.contrib.layers.xavier_initializer()):

self._init_inputs()

self.keep_prob = tf.get_variable("keep_prob", [], initializer=tf.constant_initializer(keep_prob))

with tf.device("/cpu:0"):

# embeddings

self.output_embedding = tf.get_variable("E_candidate", [answer_vocab_size, self._size],

initializer=self._init)

self.input_embedding = tf.get_variable("E_words", [vocab_size, self._size],

initializer=self._init)

answer, _ = tf.dynamic_partition(self._answer_input, self._query_partition, 2)

lookup_individual = tf.nn.embedding_lookup(self.output_embedding, answer)

cands, _ = tf.dynamic_partition(self._answer_candidates, self._query_partition, 2)

self.candidate_lookup = tf.nn.embedding_lookup(self.output_embedding, cands)

self.num_hops = tf.Variable(self._max_hops, trainable=False, name="num_queries")

self.query = self._comp_f()

answer = self._retrieve_answer(self.query)

self.score = tf_util.batch_dot(lookup_individual, answer)

self.scores_with_negs = self._score_candidates(answer)

if is_train:

self.learning_rate = tf.Variable(float(learning_rate), trainable=False, name="lr")

self.global_step = tf.Variable(0, trainable=False, name="step")

self.opt = tf.train.AdamOptimizer(self.learning_rate)

current_batch_size = tf.gather(tf.shape(self.scores_with_negs), [0])

loss = math_ops.reduce_sum(

tf.nn.sparse_softmax_cross_entropy_with_logits(self.scores_with_negs,

tf.tile(tf.constant([0], tf.int64),

current_batch_size)))

train_params = tf.trainable_variables()

self.training_weight = tf.Variable(1.0, trainable=False, name="training_weight")

self._loss = loss / math_ops.cast(current_batch_size, tf.float32)

self._grads = tf.gradients(self._loss, train_params, self.training_weight, colocate_gradients_with_ops=True)

if len(train_params) > 0:

grads, _ = tf.clip_by_global_norm(self._grads, 5.0)

self._update = self.opt.apply_gradients(zip(grads, train_params),

global_step=self.global_step)

else:

self._update = tf.assign_add(self.global_step, 1)

self.saver = tf.train.Saver(tf.all_variables(), max_to_keep=1)

def _score_candidates(self, answer):

return tf.squeeze(tf.batch_matmul(self.candidate_lookup, tf.expand_dims(answer, [2])), [2]) + \

self._candidate_mask # number of negative candidates can vary for each example

def _composition_function(self, inputs, length, init_state=None):

if self._composition == "GRU":

cell = GRUCell(self._size)

return dynamic_rnn(cell, inputs, sequence_length=length, time_major=True,

initial_state=init_state, dtype=tf.float32)[0]

elif self._composition == "LSTM":

cell = BasicLSTMCell(self._size)

init_state = tf.concat(1, [tf.zeros_like(init_state, tf.float32), init_state]) if init_state else None

outs = dynamic_rnn(cell, inputs, sequence_length=length, time_major=True,

initial_state=init_state, dtype=tf.float32)[0]

return outs

elif self._composition == "BiGRU":

cell = GRUCell(self._size // 2, self._size)

init_state_fw, init_state_bw = tf.split(1, 2, init_state) if init_state else (None, None)

with tf.variable_scope("forward"):

fw_outs = dynamic_rnn(cell, inputs, sequence_length=length, time_major=True,

initial_state=init_state_fw, dtype=tf.float32)[0]

with tf.variable_scope("backward"):

rev_inputs = tf.reverse_sequence(tf.pack(inputs), length, 0, 1)

rev_inputs = [tf.reshape(x, [-1, self._size]) for x in tf.split(0, len(inputs), rev_inputs)]

bw_outs = dynamic_rnn(cell, rev_inputs, sequence_length=length, time_major=True,

initial_state=init_state_bw, dtype=tf.float32)[0]

bw_outs = tf.reverse_sequence(tf.pack(bw_outs), length, 0, 1)

bw_outs = [tf.reshape(x, [-1, self._size]) for x in tf.split(0, len(inputs), bw_outs)]

return [tf.concat(1, [fw_out, bw_out]) for fw_out, bw_out in zip(fw_outs, bw_outs)]

else:

raise NotImplementedError("Other compositions not implemented yet.")

def name(self):

return self.__class__.__name__

def _comp_f(self):

"""

Encodes all queries (including supporting queries)

:return: encoded queries

"""

with tf.device("/cpu:0"):

max_length = tf.cast(tf.reduce_max(self._length), tf.int32)

context_t = tf.transpose(self._context)

context_t = tf.slice(context_t, [0, 0], tf.pack([max_length, -1]))

embedded = tf.nn.embedding_lookup(self.input_embedding, context_t)

embedded = tf.nn.dropout(embedded, self.keep_prob)

batch_size = tf.shape(self._context)[0]

batch_size_32 = tf.reshape(batch_size, [1])

batch_size_64 = tf.cast(batch_size, tf.int64)

with tf.device(self._device1):

#use other device for backward rnn

with tf.variable_scope("backward"):

min_end = tf.segment_min(self._ends, self._span_context)

init_state = tf.get_variable("init_state", [self._size], initializer=self._init)

init_state = tf.reshape(tf.tile(init_state, batch_size_32), [-1, self._size])

rev_embedded = tf.reverse_sequence(embedded, self._length, 0, 1)

# TIME-MAJOR: [T, B, S]

outs_bw = self._composition_function(rev_embedded, self._length - min_end, init_state)

# reshape to all possible queries for all sequences. Dim[0]=batch_size*(max_length+1).

# "+1" because we include the initial state

outs_bw = tf.reshape(tf.concat(0, [tf.expand_dims(init_state, 0), outs_bw]), [-1, self._size])

# gather respective queries via their lengths-start (because reversed sequence)

lengths_aligned = tf.gather(self._length, self._span_context)

out_bw = tf.gather(outs_bw, (lengths_aligned - self._ends) * batch_size_64 + self._span_context)

with tf.device(self._device2):

with tf.variable_scope("forward"):

#e_inputs = [tf.reshape(e, [-1, self._size]) for e in tf.split(1, self._max_length, embedded)]

max_start = tf.segment_max(self._starts, self._span_context)

init_state = tf.get_variable("init_state", [self._size], initializer=self._init)

init_state = tf.reshape(tf.tile(init_state, batch_size_32), [-1, self._size])

# TIME-MAJOR: [T, B, S]

outs_fw = self._composition_function(embedded, max_start, init_state)

# reshape to all possible queries for all sequences. Dim[0]=batch_size*(max_length+1).

# "+1" because we include the initial state

outs_fw = tf.reshape(tf.concat(0, [tf.expand_dims(init_state, 0), outs_fw]), [-1, self._size])

# gather respective queries via their positions (with offset of batch_size*ends)

out_fw = tf.gather(outs_fw, self._starts * batch_size_64 + self._span_context)

# form query from forward and backward compositions

query = tf.contrib.layers.fully_connected(tf.concat(1, [out_fw, out_bw]), self._size,

activation_fn=None, weights_initializer=None, biases_initializer=None)

query = tf.add_n([query, out_bw, out_fw])

return query

def set_train(self, sess):

"""

enables dropout

:param sess:

:return:

"""

sess.run(self.keep_prob.initializer)

def set_eval(self, sess):

"""

removes dropout

:param sess:

:return:

"""

sess.run(self.keep_prob.assign(1.0))

def _retrieve_answer(self, query):

"""

Retrieves answer based on the specified query. Implements consecutive updates to the query and answer.

:return: answer, if num_hops is 0, returns query itself

"""

query, supp_queries = tf.dynamic_partition(query, self._query_partition, 2)

with tf.variable_scope("support"):

num_queries = tf.shape(query)[0]

with tf.device("/cpu:0"):

_, supp_answer_output_ids = tf.dynamic_partition(self._answer_input, self._query_partition, 2)

_, supp_answer_input_ids = tf.dynamic_partition(self._answer_word_input, self._query_partition, 2)

supp_answers = tf.nn.embedding_lookup(self.output_embedding, supp_answer_output_ids)

aligned_supp_answers = tf.gather(supp_answers, self._support_ids) # and with respective answers

if self._max_hops > 1:

# used in multihop

answer_words = tf.nn.embedding_lookup(self.input_embedding, supp_answer_input_ids)

aligned_answers_input = tf.gather(answer_words, self._support_ids)

self.support_scores = []

query_as_answer = tf.contrib.layers.fully_connected(query, self._size,

activation_fn=None, weights_initializer=None,

biases_initializer=None, scope="query_to_answer")

query_as_answer = query_as_answer * tf.sigmoid(tf.get_variable("query_as_answer_gate", tuple(),

initializer=tf.constant_initializer(0.0)))

current_answer = query_as_answer

current_query = query

aligned_support = tf.gather(supp_queries, self._support_ids) # align supp_queries with queries

collab_support = tf.gather(query, self._collab_support_ids) # align supp_queries with queries

aligned_support = tf.concat(0, [aligned_support, collab_support])

query_ids = tf.concat(0, [self._query_ids, self._collab_query_ids])

self.answer_weights = []

for i in range(self._max_hops):

if i > 0:

tf.get_variable_scope().reuse_variables()

collab_queries = tf.gather(current_query, self._collab_query_ids) # align supp_queries with queries

aligned_queries = tf.gather(current_query, self._query_ids) # align queries

aligned_queries = tf.concat(0, [aligned_queries, collab_queries])

with tf.variable_scope("support_scores"):

scores = tf_util.batch_dot(aligned_queries, aligned_support)

self.support_scores.append(scores)

score_max = tf.gather(tf.segment_max(scores, query_ids), query_ids)

e_scores = tf.exp(scores - score_max)

norm = tf.unsorted_segment_sum(e_scores, query_ids, num_queries) + 0.00001 # for zero norms

norm = tf.expand_dims(norm, 1)

e_scores = tf.expand_dims(e_scores, 1)

with tf.variable_scope("support_answers"):

aligned_supp_answers_with_collab = tf.concat(0, [aligned_supp_answers, collab_queries])

weighted_supp_answers = tf.unsorted_segment_sum(e_scores * aligned_supp_answers_with_collab,

query_ids, num_queries) / norm

with tf.variable_scope("support_queries"):

weighted_supp_queries = tf.unsorted_segment_sum(e_scores * aligned_support, query_ids, num_queries) / norm

with tf.variable_scope("answer_accumulation"):

answer_p_max = tf.reduce_max(tf.nn.softmax(self._score_candidates(weighted_supp_answers)), [1], keep_dims=True)

answer_weight = tf.contrib.layers.fully_connected(tf.concat(1, [query_as_answer * weighted_supp_answers,

weighted_supp_queries * current_query,

answer_p_max]),

1,

activation_fn=tf.nn.sigmoid,

weights_initializer=tf.constant_initializer(0.0),

biases_initializer=tf.constant_initializer(0.0),

scope="answer_weight")

new_answer = answer_weight * weighted_supp_answers + current_answer

# this condition allows for setting varying number of hops

current_answer = tf.cond(tf.greater(self.num_hops, i),

lambda: new_answer,

lambda: current_answer)

self.answer_weights.append(answer_weight)

if i < self._max_hops - 1:

with tf.variable_scope("query_update"):

# prepare subsequent query

aligned_answers_input_with_collab = tf.concat(0, [aligned_answers_input, collab_queries])

weighted_answer_words = tf.unsorted_segment_sum(e_scores * aligned_answers_input_with_collab,

query_ids, num_queries) / norm

c = tf.contrib.layers.fully_connected(tf.concat(1, [current_query, weighted_supp_queries, weighted_answer_words]),

self._size, activation_fn=tf.tanh, scope="update_candidate",

weights_initializer=None, biases_initializer=None)

gate = tf.contrib.layers.fully_connected(tf.concat(1, [current_query, weighted_supp_queries]),

self._size, activation_fn=tf.sigmoid,

weights_initializer=None, scope="update_gate",

biases_initializer=tf.constant_initializer(1))

current_query = gate * current_query + (1-gate) * c

return current_answer

def _init_inputs(self):

#General

with tf.device("/cpu:0"):

self._context = tf.placeholder(tf.int64, shape=[None, self._max_length], name="context")

self._answer_candidates = tf.placeholder(tf.int64, shape=[None, None], name="candidates")

self._answer_input = tf.placeholder(tf.int64, shape=[None], name="answer")

# answer word ids (index to E_embeddings) might differ from answer ids (input to E_candidates)

self._answer_word_input = tf.placeholder(tf.int64, shape=[None], name="answer_word")

self._starts = tf.placeholder(tf.int64, shape=[None], name="span_start")

self._ends = tf.placeholder(tf.int64, shape=[None], name="span_end")

# holds batch idx for respective span

self._span_context = tf.placeholder(tf.int64, shape=[None], name="answer_position_context")

self._candidate_mask = tf.placeholder(tf.float32, shape=[None, None], name="candidate_mask")

self._length = tf.placeholder(tf.int64, shape=[None], name="context_length")

self._ctxt = np.zeros([self._batch_size, self._max_length], dtype=np.int64)

self._len = np.zeros([self._batch_size], dtype=np.int64)

#Supporting Evidence

# partition of queries (class 0) and support (class 1)

self._query_partition = tf.placeholder(tf.int32, [None], "query_partition")

# aligned support ids with query ids for supporting evidence

self._support_ids = tf.placeholder(tf.int64, shape=[None], name="support_for_query_ids")

self._collab_support_ids = tf.placeholder(tf.int64, shape=[None], name="collab_supp_ids")

self._query_ids = tf.placeholder(tf.int64, shape=[None], name="query_for_support_ids")

self._collab_query_ids = tf.placeholder(tf.int64, shape=[None], name="collab_query_ids")

self._feed_dict = {}

def _change_batch_size(self, batch_size):

new_ctxt_in = np.zeros([batch_size, self._max_length], dtype=np.int64)

new_ctxt_in[:self._batch_size] = self._ctxt

self._ctxt = new_ctxt_in

new_length = np.zeros([batch_size], dtype=np.int64)

new_length[:self._batch_size] = self._len

self._len = new_length

self._batch_size = batch_size

def _start_adding_examples(self):

self._batch_idx = 0

self._query_idx = 0

self._support_idx = 0

self._answer_cands = []

self._answer_in = []

self._answer_word_in = []

self._s = []

self._e = []

self._span_ctxt = []

# supporting evidence

self._query_part = []

self.queries_for_support = []

self.support_for_queries = []

self._collab_queries = []

self._collab_support = []

self.supporting_qa = []

def _add_example(self, context_queries, is_query=True):

'''

All queries and supporting queries are encoded the same. However we keep track of which are queries,

which are support and how they belong to each other via partition variables and aligned support to query

and query to support ids.

:param context_queries: contains all queries about a particular context, see ContextQueries in query.py

:param is_query: True if this is query, False if this is support

:return:

'''

assert is_query or context_queries.support is None, "Support cannot have support!"

if self._batch_idx >= self._batch_size:

self._change_batch_size(max(self._batch_size*2, self._batch_idx))

self._ctxt[self._batch_idx][:len(context_queries.context)] = context_queries.context

self._len[self._batch_idx] = len(context_queries.context)

batch_idx = self._batch_idx

self._batch_idx += 1

for i, q in enumerate(context_queries.queries):

self._s.append(q.start)

self._e.append(q.end)

self._span_ctxt.append(batch_idx)

self._answer_in.append(q.answer if q.answer is not None else q.candidates[0])

self._answer_word_in.append(q.answer_word)

cands = [q.answer] if q.answer is not None else []

if q.candidates is not None:

cands.extend(c for c in q.candidates if c != q.answer)

self._answer_cands.append(cands)

self._query_part.append(0 if is_query else 1)

if is_query:

if context_queries.collaborative_support:

# save queries also as support, only with different query_partition index (1 for support)

for i in range(len(context_queries.queries)):

for j in range(len(context_queries.queries)):

if j != i:

self._collab_queries.append(self._query_idx+i)

self._collab_support.append(self._query_idx+j)

### add query specific supports ###

for i, q in enumerate(context_queries.queries):

if q.support is not None and self._max_hops > 0:

for qs in q.support:

if qs.context is None:

#supporting context is the same as query context, only add corresponding positions

for q in qs.queries:

self._s.append(q.start)

self._e.append(q.end)

self._span_ctxt.append(batch_idx)

self._answer_in.append(q.answer)

self._answer_word_in.append(q.answer_word)

self._answer_cands.append([q.answer])

self._query_part.append(1)

self.supporting_qa.append((q.context, q.start, q.end, q.answer))

else:

self._add_example(qs, is_query=False)

# align queries with support idxs

self.support_for_queries.extend(range(self._support_idx, self._support_idx+len(qs.queries)))

self._support_idx += len(qs.queries)

self.queries_for_support.extend([self._query_idx] * len(qs.queries))

self._query_idx += 1

### add context specific support to all queries of this context ###

if context_queries.support is not None and self._max_hops > 0:

for qs in context_queries.support:

if qs.context is None:

for q in qs.queries:

self._s.append(q.start)

self._e.append(q.end)

self._span_ctxt.append(batch_idx)

self._answer_in.append(q.answer)

self._answer_word_in.append(q.answer_word)

self._answer_cands.append([q.answer])

self._query_part.append(1)

self.supporting_qa.append((q.context, q.start, q.end, q.answer))

else:

self._add_example(qs, is_query=False)

# this evidence supports all queries in this context

for i, _ in enumerate(context_queries.queries):

# align queries with support idxs

self.support_for_queries.extend(range(self._support_idx, self._support_idx+len(qs.queries)))

self.queries_for_support.extend([self._query_idx - len(context_queries.queries) + i] * len(qs.queries))

self._support_idx += len(qs.queries)

else:

for i, q in enumerate(context_queries.queries):

self.supporting_qa.append((q.context, q.start, q.end, q.answer))

def _finish_adding_examples(self):

max_cands = max((len(x) for x in self._answer_cands))

# mask is used to determine which candidates are real candidates and which are dummies,

# number of candidates can vary from query to query within a batch

cand_mask = []

for i in range(len(self._answer_cands)):

l = len(self._answer_cands[i])

if self._query_part[i] == 0: # if this is a query (and not supporting evidence)

mask = [0] * l

for _ in range(max_cands - l):

self._answer_cands[i].append(self._answer_cands[i][0]) # dummy

if self._query_part[i] == 0:

mask.append(-1e6) # this is added to scores, serves basically as a bias mask to exclude dummy negative candidates

if self._query_part[i] == 0:

cand_mask.append(mask)

if self._batch_idx < self._batch_size:

self._feed_dict[self._context] = self._ctxt[:self._batch_idx]

self._feed_dict[self._length] = self._len[:self._batch_idx]

else:

self._feed_dict[self._context] = self._ctxt

self._feed_dict[self._length] = self._len

self._feed_dict[self._starts] = self._s

self._feed_dict[self._ends] = self._e

self._feed_dict[self._span_context] = self._span_ctxt

self._feed_dict[self._answer_input] = self._answer_in

self._feed_dict[self._answer_word_input] = self._answer_word_in

self._feed_dict[self._answer_candidates] = self._answer_cands

self._feed_dict[self._candidate_mask] = cand_mask

self._feed_dict[self._query_ids] = self.queries_for_support

self._feed_dict[self._support_ids] = self.support_for_queries

self._feed_dict[self._collab_query_ids] = self._collab_queries

self._feed_dict[self._collab_support_ids] = self._collab_support

self._feed_dict[self._query_partition] = self._query_part

def get_feed_dict(self):

return self._feed_dict

def step(self, sess, queries, mode="update"):

'''

:param sess:

:param queries: list of ContextQueries

:param mode: "loss" for loss, else performs update on parameters

:return:

'''

assert self._is_train, "model has to be created in training mode!"

if mode == "loss":

return self.run(sess, self._loss, queries)

else:

return self.run(sess, [self._loss, self._update], queries)[0]

def run(self, sess, to_run, queries):

'''

:param sess:

:param to_run: target(s) to run, e.g. :

* self.num_hops,

* self.query,

* self.score (only score for provided answers),

* self.score_with_negs (scores of all candidates where score[0] is score of the answer)

* self.input_embedding,

* self.output_embedding,

* self.support_scores (match-scores for all support with query for each hop. Aligns with

self.supporting_qa which keeps track of all support QA-pairs and self.queries_for_support which

defines batch_idx of query for all supporting_qa),,

* self.answer_weights (weight used to accumulate retrieved answer in each hop)

answer = self._retrieve_answer(self.query)

self.score = tf_util.batch_dot(lookup_individual, answer)

self.scores_with_negs = self._score_candidates(answer)

if is_train:

self.learning_rate = tf.Variable(float(learning_rate), trainable=False, name="lr")

self.global_step = tf.Variable(0, trainable=False, name="step")

:param queries: list of ContextQueries

:return:

'''

batch_size = len(queries)

self._start_adding_examples()

num_batch_queries = 0

for batch_idx in range(batch_size):

context_query = queries[batch_idx]

num_batch_queries += len(context_query.queries)

self._add_example(context_query)

self._finish_adding_examples()

model = QAModel(10, 4, 5, 5, 5, max_hops=2)

# 3 contexts (of length 3) with queries at 2/1/2 (totaling 5) positions

# and respective negative candidates for each position

contexts = [[0, 1, 2] , [1, 2, 0], [0, 2, 1]] # 4 => placeholder for prediction position

support = [ContextQueries(contexts[0], [ClozeQuery(contexts[0], 0,1,0,0,[2,1]),

ClozeQuery(contexts[0], 2,3,2,2,[0,1])]),

ContextQueries(contexts[1], [ClozeQuery(contexts[1], 1,2,2,2,[0,1])])]

queries = [ContextQueries(contexts[0], [ClozeQuery(contexts[0], 0,1,0,0,[2,1], support=support),

ClozeQuery(contexts[0], 2,3,2,2,[0,1], support=support)]),

ContextQueries(contexts[1], [ClozeQuery(contexts[1], 1,2,2,2,[0,1], support=support)]),

ContextQueries(contexts[2], [ClozeQuery(contexts[2], 1, 2, 1, 1, [0,2], support=support),

ClozeQuery(contexts[2], 2,3,2,2,[0,1], support=support)])]

with tf.Session() as sess:

sess.run(tf.initialize_all_variables())

sess.run(model.num_hops.assign(1))

print("Test update ...")

for i in range(10):

print("Loss: %.3f" %

model.step(sess, queries)[0])

print("Test scoring ...")

print(model.run(sess, model.scores_with_negs, queries))