"""
Pystruct-compatible models.
"""
# Author: Vlad Niculae <vlad@vene.ro>
# License: BSD 3-clause
# AD3 is (c) Andre F. T. Martins, LGPLv3.0: http://www.cs.cmu.edu/~ark/AD3/
import warnings
import numpy as np
from sklearn.utils import compute_class_weight
from sklearn.utils.extmath import safe_sparse_dot
from sklearn.metrics import f1_score
from sklearn.preprocessing import LabelEncoder, label_binarize
from pystruct.models import StructuredModel
from marseille.inference import loss_augment_unaries, CDCP_ILLEGAL_LINKS
from marseille.argdoc import DocLabel
from marseille.custom_logging import logging
from itertools import permutations
from ad3 import factor_graph as fg
def _binary_2d(y):
if y.shape[1] == 1:
y = np.column_stack([1 - y, y])
return y
def arg_f1_scores(Y_true, Y_pred, **kwargs):
macro = []
micro_true = []
micro_pred = []
for y_true, y_pred in zip(Y_true, Y_pred):
macro.append(f1_score(y_true, y_pred, **kwargs))
micro_true.extend(y_true)
micro_pred.extend(y_pred)
return np.mean(macro), f1_score(micro_true, micro_pred, **kwargs)
class BaseArgumentMixin(object):
def initialize_labels(self, Y):
y_nodes_flat = [y_val for y in Y for y_val in y.nodes]
y_links_flat = [y_val for y in Y for y_val in y.links]
self.prop_encoder_ = LabelEncoder().fit(y_nodes_flat)
self.link_encoder_ = LabelEncoder().fit(y_links_flat)
self.n_prop_states = len(self.prop_encoder_.classes_)
self.n_link_states = len(self.link_encoder_.classes_)
self.prop_cw_ = np.ones_like(self.prop_encoder_.classes_,
dtype=np.double)
self.link_cw_ = compute_class_weight(self.class_weight,
self.link_encoder_.classes_,
y_links_flat)
self.link_cw_ /= self.link_cw_.min()
logging.info('Setting node class weights {}'.format(", ".join(
"{}: {}".format(lbl, cw) for lbl, cw in zip(
self.prop_encoder_.classes_, self.prop_cw_))))
logging.info('Setting link class weights {}'.format(", ".join(
"{}: {}".format(lbl, cw) for lbl, cw in zip(
self.link_encoder_.classes_, self.link_cw_))))
def _round(self, prop_marg, link_marg, prop_unary=None, link_unary=None,
inverse_transform=True):
# ensure ties are broken according to unary scores
if prop_unary is not None:
prop_unary = prop_unary.copy()
prop_unary -= np.min(prop_unary)
prop_unary /= np.max(prop_unary) * np.max(prop_marg)
prop_marg[prop_marg > 1e-9] += prop_unary[prop_marg > 1e-9]
if link_unary is not None:
link_unary = link_unary.copy()
link_unary -= np.min(link_unary)
link_unary /= np.max(link_unary) * np.max(link_marg)
link_marg[link_marg > 1e-9] += link_unary[link_marg > 1e-9]
y_hat_props = np.argmax(prop_marg, axis=1)
y_hat_links = np.argmax(link_marg, axis=1)
if inverse_transform:
y_hat_props = self.prop_encoder_.inverse_transform(y_hat_props)
y_hat_links = self.link_encoder_.inverse_transform(y_hat_links)
return DocLabel(y_hat_props, y_hat_links)
def loss(self, y, y_hat):
if not isinstance(y_hat, DocLabel):
return self.continuous_loss(y, y_hat)
y_nodes = self.prop_encoder_.transform(y.nodes)
y_links = self.link_encoder_.transform(y.links)
node_loss = np.sum(self.prop_cw_[y_nodes] * (y.nodes != y_hat.nodes))
link_loss = np.sum(self.link_cw_[y_links] * (y.links != y_hat.links))
return node_loss + link_loss
def max_loss(self, y):
y_nodes = self.prop_encoder_.transform(y.nodes)
y_links = self.link_encoder_.transform(y.links)
return np.sum(self.prop_cw_[y_nodes]) + np.sum(self.link_cw_[y_links])
def continuous_loss(self, y, y_hat):
if isinstance(y_hat, DocLabel):
raise ValueError("continuous loss on discrete input")
if isinstance(y_hat[0], tuple):
y_hat = y_hat[0]
prop_marg, link_marg = y_hat
y_nodes = self.prop_encoder_.transform(y.nodes)
y_links = self.link_encoder_.transform(y.links)
prop_ix = np.indices(y.nodes.shape)
link_ix = np.indices(y.links.shape)
# relies on prop_marg and link_marg summing to 1 row-wise
prop_loss = np.sum(self.prop_cw_[y_nodes] *
(1 - prop_marg[prop_ix, y_nodes]))
link_loss = np.sum(self.link_cw_[y_links] *
(1 - link_marg[link_ix, y_links]))
loss = prop_loss + link_loss
return loss
def _marg_rounded(self, x, y):
y_node = y.nodes
y_link = y.links
Y_node = label_binarize(y_node, self.prop_encoder_.classes_)
Y_link = label_binarize(y_link, self.link_encoder_.classes_)
# XXX can this be avoided?
Y_node, Y_link = map(_binary_2d, (Y_node, Y_link))
src_type = Y_node[x.link_to_prop[:, 0]]
trg_type = Y_node[x.link_to_prop[:, 1]]
if self.compat_features:
pw = np.einsum('...j,...k,...l->...jkl',
src_type, trg_type, Y_link)
compat = np.tensordot(x.X_compat.T, pw, axes=[1, 0])
else:
# equivalent to compat_features == np.ones(n_links)
compat = np.einsum('ij,ik,il->jkl', src_type, trg_type, Y_link)
second_order = []
if self.coparents_ or self.grandparents_ or self.siblings_:
link = {(a, b): k for k, (a, b) in enumerate(x.link_to_prop)}
if self.coparents_:
second_order.extend(y_link[link[a, b]] & y_link[link[c, b]]
for a, b, c in x.second_order)
if self.grandparents_:
second_order.extend(y_link[link[a, b]] & y_link[link[b, c]]
for a, b, c in x.second_order)
if self.siblings_:
second_order.extend(y_link[link[b, a]] & y_link[link[b, c]]
for a, b, c in x.second_order)
second_order = np.array(second_order)
return Y_node, Y_link, compat, second_order
def _marg_fractional(self, x, y):
(prop_marg, link_marg), (compat_marg, second_order_marg) = y
if self.compat_features:
compat_marg = np.tensordot(x.X_compat.T, compat_marg, axes=[1, 0])
else:
compat_marg = compat_marg.sum(axis=0)
return prop_marg, link_marg, compat_marg, second_order_marg
def _inference(self, x, potentials, exact=False, relaxed=True,
return_energy=False, constraints=None,
eta=0.1, adapt=True, max_iter=5000,
verbose=False):
(prop_potentials,
link_potentials,
compat_potentials,
coparent_potentials,
grandparent_potentials,
sibling_potentials) = potentials
n_props, n_prop_classes = prop_potentials.shape
n_links, n_link_classes = link_potentials.shape
g = fg.PFactorGraph()
g.set_verbosity(verbose)
prop_vars = [g.create_multi_variable(n_prop_classes)
for _ in range(n_props)]
link_vars = [g.create_multi_variable(n_link_classes)
for _ in range(n_links)]
for var, scores in zip(prop_vars, prop_potentials):
for state, score in enumerate(scores):
var.set_log_potential(state, score)
for var, scores in zip(link_vars, link_potentials):
for state, score in enumerate(scores):
var.set_log_potential(state, score)
# compatibility trigram factors
compat_factors = []
link_vars_dict = {}
link_on, link_off = self.link_encoder_.transform([True, False])
# account for compat features
if self.compat_features:
assert compat_potentials.shape[0] == n_links
compats = compat_potentials
else:
compats = (compat_potentials for _ in range(n_links))
for (src, trg), link_v, compat in zip(x.link_to_prop,
link_vars,
compats):
src_v = prop_vars[src]
trg_v = prop_vars[trg]
compat_factors.append(g.create_factor_dense([src_v, trg_v, link_v],
compat.ravel()))
# keep track of binary link variables, for constraints.
# we need .get_state() to get the underlaying PBinaryVariable
link_vars_dict[src, trg] = link_v.get_state(link_on)
# second-order factors
coparent_factors = []
grandparent_factors = []
sibling_factors = []
for score, (a, b, c) in zip(coparent_potentials, x.second_order):
# a -> b <- c
vars = [link_vars_dict[a, b], link_vars_dict[c, b]]
coparent_factors.append(g.create_factor_pair(vars, score))
for score, (a, b, c) in zip(grandparent_potentials, x.second_order):
# a -> b -> c
vars = [link_vars_dict[a, b], link_vars_dict[b, c]]
grandparent_factors.append(g.create_factor_pair(vars, score))
for score, (a, b, c) in zip(sibling_potentials, x.second_order):
# a <- b -> c
vars = [link_vars_dict[b, a], link_vars_dict[b, c]]
sibling_factors.append(g.create_factor_pair(vars, score))
# domain-specific constraints
if constraints and 'cdcp' in constraints:
# antisymmetry: if a -> b, then b cannot -> a
for src in range(n_props):
for trg in range(src):
fwd_link_v = link_vars_dict[src, trg]
rev_link_v = link_vars_dict[trg, src]
g.create_factor_logic('ATMOSTONE',
[fwd_link_v, rev_link_v],
[False, False])
# transitivity.
# forall a != b != c: a->b and b->c imply a->c
for a, b, c in permutations(range(n_props), 3):
ab_link_v = link_vars_dict[a, b]
bc_link_v = link_vars_dict[b, c]
ac_link_v = link_vars_dict[a, c]
g.create_factor_logic('IMPLY',
[ab_link_v, bc_link_v, ac_link_v],
[False, False, False])
# standard model:
if 'strict' in constraints:
for src, trg in x.link_to_prop:
src_v = prop_vars[src]
trg_v = prop_vars[trg]
for types in CDCP_ILLEGAL_LINKS:
src_ix, trg_ix = self.prop_encoder_.transform(types)
g.create_factor_logic('IMPLY',
[src_v.get_state(src_ix),
trg_v.get_state(trg_ix),
link_vars_dict[src, trg]],
[False, False, True])
elif constraints and 'ukp' in constraints:
# Tree constraints using AD3 MST factor for each paragraph.
# First, identify paragraphs
prop_para = np.array(x.prop_para)
link_para = prop_para[x.link_to_prop[:, 0]]
tree_factors = []
for para_ix in np.unique(link_para):
props = np.where(prop_para == para_ix)[0]
offset = props.min()
para_vars = []
para_arcs = [] # call them arcs, semantics differ from links
# add a new head node pointing to every possible variable
for relative_ix, prop_ix in enumerate(props, 1):
para_vars.append(g.create_binary_variable())
para_arcs.append((0, relative_ix))
# add an MST arc for each link
for src, trg in x.link_to_prop[link_para == para_ix]:
relative_src = src - offset + 1
relative_trg = trg - offset + 1
para_vars.append(link_vars_dict[src, trg])
# MST arcs have opposite direction from argument links!
# because each prop can have multiple supports but not
# the other way around
para_arcs.append((relative_trg, relative_src))
tree = fg.PFactorTree()
g.declare_factor(tree, para_vars, True)
tree.initialize(1 + len(props), para_arcs)
tree_factors.append(tree)
if 'strict' in constraints:
# further domain-specific constraints
mclaim_ix, claim_ix, premise_ix = self.prop_encoder_.transform(
['MajorClaim', 'Claim', 'Premise'])
# a -> b implies a = 'premise'
for (src, trg), link_v in zip(x.link_to_prop, link_vars):
src_v = prop_vars[src]
g.create_factor_logic('IMPLY',
[link_v.get_state(link_on),
src_v.get_state(premise_ix)],
[False, False])
g.fix_multi_variables_without_factors()
g.set_eta_ad3(eta)
g.adapt_eta_ad3(adapt)
g.set_max_iterations_ad3(max_iter)
if exact:
val, posteriors, additionals, status = g.solve_exact_map_ad3()
else:
val, posteriors, additionals, status = g.solve_lp_map_ad3()
status = ["integer", "fractional", "infeasible", "not solved"][status]
prop_marg = posteriors[:n_props * n_prop_classes]
prop_marg = np.array(prop_marg).reshape(n_props, -1)
link_marg = posteriors[n_props * n_prop_classes:]
# remaining posteriors are for artificial root nodes for MST factors
link_marg = link_marg[:n_links * n_link_classes]
link_marg = np.array(link_marg).reshape(n_links, -1)
n_compat = n_links * n_link_classes * n_prop_classes ** 2
compat_marg = additionals[:n_compat]
compat_marg = np.array(compat_marg).reshape((n_links,
n_prop_classes,
n_prop_classes,
n_link_classes))
second_ordermarg = np.array(additionals[n_compat:])
posteriors = (prop_marg, link_marg)
additionals = (compat_marg, second_ordermarg)
if relaxed:
y_hat = posteriors, additionals
else:
y_hat = self._round(prop_marg, link_marg, prop_potentials,
link_potentials)
if return_energy:
return y_hat, status, -val
else:
return y_hat, status
def _score(self, Y_true, Y_pred):
acc = sum(1 for y_true, y_pred in zip(Y_true, Y_pred)
if np.all(y_true.links == y_pred.links) and
np.all(y_true.nodes == y_pred.nodes))
acc /= len(Y_true)
with warnings.catch_warnings():
warnings.simplefilter('ignore')
link_macro, link_micro = arg_f1_scores(
(y.links for y in Y_true),
(y.links for y in Y_pred),
average='binary',
pos_label=True,
labels=self.link_encoder_.classes_
)
node_macro, node_micro = arg_f1_scores(
(y.nodes for y in Y_true),
(y.nodes for y in Y_pred),
average='macro',
labels=self.prop_encoder_.classes_
)
return link_macro, link_micro, node_macro, node_micro, acc
class ArgumentGraphCRF(BaseArgumentMixin, StructuredModel):
def __init__(self, class_weight=None, link_node_weight_ratio=1,
exact=False, constraints=None, compat_features=False,
coparents=False, grandparents=False, siblings=False):
self.class_weight = class_weight
self.link_node_weight_ratio = link_node_weight_ratio
self.exact = exact
self.constraints = constraints
self.compat_features = compat_features
self.coparents = coparents
self.grandparents = grandparents
self.siblings = siblings
self.n_second_order_factors_ = coparents + grandparents + siblings
self.n_prop_states = None
self.n_link_states = None
self.n_prop_features = None
self.n_link_features = None
self.n_second_order_features_ = None
self.n_compat_features_ = None
self.inference_calls = 0
super(ArgumentGraphCRF, self).__init__()
def initialize(self, X, Y):
# each x in X is a vectorized doc exposing sp.csr x.X_prop, x.X_link,
# and maybe x.X_compat and x.X_sec_ord
# each y in Y exposes lists y.nodes, y.links
x = X[0]
self.n_prop_features = x.X_prop.shape[1]
self.n_link_features = x.X_link.shape[1]
if self.compat_features:
self.n_compat_features_ = x.X_compat.shape[1]
if self.n_second_order_factors_:
self.n_second_order_features_ = x.X_sec_ord.shape[1]
else:
self.n_second_order_features_ = 0
self.initialize_labels(Y)
self._set_size_joint_feature()
self.coparents_ = self.coparents
self.grandparents_ = self.grandparents
self.siblings_ = self.siblings
def _set_size_joint_feature(self): # assumes no second order
compat_size = self.n_prop_states ** 2 * self.n_link_states
if self.compat_features:
compat_size *= self.n_compat_features_
total_n_second_order = (self.n_second_order_features_ *
self.n_second_order_factors_)
self.size_joint_feature = (self.n_prop_features * self.n_prop_states +
self.n_link_features * self.n_link_states +
compat_size + total_n_second_order)
logging.info("Joint feature size: {}".format(self.size_joint_feature))
def joint_feature(self, x, y):
if isinstance(y, DocLabel):
Y_prop, Y_link, compat, second_order = self._marg_rounded(x, y)
else:
Y_prop, Y_link, compat, second_order = self._marg_fractional(x, y)
prop_acc = safe_sparse_dot(Y_prop.T, x.X_prop) # node_cls * node_feats
link_acc = safe_sparse_dot(Y_link.T, x.X_link) # link_cls * link_feats
f_sec_ord = []
if len(second_order):
second_order = second_order.reshape(-1, len(x.second_order))
if self.coparents:
f_sec_ord.append(safe_sparse_dot(second_order[0], x.X_sec_ord))
second_order = second_order[1:]
if self.grandparents:
f_sec_ord.append(safe_sparse_dot(second_order[0], x.X_sec_ord))
second_order = second_order[1:]
if self.siblings:
f_sec_ord.append(safe_sparse_dot(second_order[0], x.X_sec_ord))
elif self.n_second_order_factors_:
# document has no second order factors so the joint feature
# must be filled with zeros manually
f_sec_ord = [np.zeros(self.n_second_order_features_)
for _ in range(self.n_second_order_factors_)]
jf = np.concatenate([prop_acc.ravel(), link_acc.ravel(),
compat.ravel()] + f_sec_ord)
return jf
# basically reversing the joint feature
def _get_potentials(self, x, w):
# check sizes?
n_node_coefs = self.n_prop_states * self.n_prop_features
n_link_coefs = self.n_link_states * self.n_link_features
n_compat_coefs = self.n_prop_states ** 2 * self.n_link_states
if self.compat_features:
n_compat_coefs *= self.n_compat_features_
assert w.size == (n_node_coefs + n_link_coefs + n_compat_coefs +
self.n_second_order_features_ *
self.n_second_order_factors_)
w_node = w[:n_node_coefs]
w_node = w_node.reshape(self.n_prop_states, self.n_prop_features)
w_link = w[n_node_coefs:n_node_coefs + n_link_coefs]
w_link = w_link.reshape(self.n_link_states, self.n_link_features)
# for readability, consume w. This is not inplace, don't worry.
w = w[n_node_coefs + n_link_coefs:]
w_compat = w[:n_compat_coefs]
if self.compat_features:
w_compat = w_compat.reshape((self.n_compat_features_, -1))
w_compat = np.dot(x.X_compat, w_compat)
compat_potentials = w_compat.reshape((-1,
self.n_prop_states,
self.n_prop_states,
self.n_link_states))
else:
compat_potentials = w_compat.reshape(self.n_prop_states,
self.n_prop_states,
self.n_link_states)
w = w[n_compat_coefs:]
coparent_potentials = grandparent_potentials = sibling_potentials = []
if self.coparents:
w_coparent = w[:self.n_second_order_features_]
coparent_potentials = safe_sparse_dot(x.X_sec_ord, w_coparent)
w = w[self.n_second_order_features_:]
if self.grandparents:
w_grandparent = w[:self.n_second_order_features_]
grandparent_potentials = safe_sparse_dot(x.X_sec_ord,
w_grandparent)
w = w[self.n_second_order_features_:]
if self.siblings:
w_sibling = w[:self.n_second_order_features_]
sibling_potentials = safe_sparse_dot(x.X_sec_ord, w_sibling)
prop_potentials = safe_sparse_dot(x.X_prop, w_node.T)
link_potentials = safe_sparse_dot(x.X_link, w_link.T)
return (prop_potentials, link_potentials, compat_potentials,
coparent_potentials, grandparent_potentials,
sibling_potentials)
def inference(self, x, w, relaxed=False, return_energy=False):
self.inference_calls += 1
potentials = self._get_potentials(x, w)
out = self._inference(x, potentials, exact=self.exact,
relaxed=relaxed, return_energy=return_energy,
constraints=self.constraints)
if return_energy:
return out[0], out[-1]
else:
return out[0]
def loss_augmented_inference(self, x, y, w, relaxed=None):
self.inference_calls += 1
potentials = self._get_potentials(x, w)
(prop_potentials,
link_potentials,
compat_potentials,
coparent_potentials,
grandparent_potentials,
sibling_potentials) = potentials
y_prop = self.prop_encoder_.transform(y.nodes)
y_link = self.link_encoder_.transform(y.links)
loss_augment_unaries(prop_potentials, y_prop, self.prop_cw_)
loss_augment_unaries(link_potentials, y_link, self.link_cw_)
potentials = (prop_potentials,
link_potentials,
compat_potentials,
coparent_potentials,
grandparent_potentials,
sibling_potentials)
out = self._inference(x, potentials, exact=self.exact,
relaxed=relaxed, constraints=self.constraints)
return out[0]
