TreeCRF.py

#!/usr/bin/env python3

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import itertools
import utils
import random
  
class ConstituencyTreeCRF(nn.Module):
  def __init__(self):
    super(ConstituencyTreeCRF, self).__init__()
    self.huge = 1e9

  def logadd(self, x, y):
    d = torch.max(x,y)  
    return torch.log(torch.exp(x-d) + torch.exp(y-d)) + d    

  def logsumexp(self, x, dim=1):
    d = torch.max(x, dim)[0]
    return torch.log(torch.exp(x - d.unsqueeze(dim).expand_as(x)).sum(dim)) + d

  def _init_table(self, scores):
    # initialize dynamic programming table
    batch_size = scores.size(0)
    n = scores.size(1)
    self.alpha = [[scores.new(batch_size).fill_(-self.huge) for _ in range(n)] for _ in range(n)]

  def _forward(self, scores):
    #inside step
    batch_size = scores.size(0)
    n = scores.size(1)
    self._init_table(scores)
    for i in range(n):
      self.alpha[i][i] = scores[:, i, i]
    for k in np.arange(1, n+1):
      for s in range(n):
        t = s + k
        if t > n-1:
          break
        tmp = [self.alpha[s][u] + self.alpha[u+1][t] + scores[:, s, t] for u in np.arange(s,t)]
        tmp = torch.stack(tmp, 1)
        self.alpha[s][t] = self.logsumexp(tmp, 1)
            
  def _backward(self, scores):
    #outside step
    batch_size = scores.size(0)
    n = scores.size(1)
    self.beta = [[None for _ in range(n)] for _ in range(n)]
    self.beta[0][n-1] = scores.new(batch_size).fill_(0)
    for k in np.arange(n-1, 0, -1):
      for s in range(n):
        t = s + k
        if t > n-1:
          break
        for u in np.arange(s, t):                    
          if s < u+1:
            tmp = self.beta[s][t] + self.alpha[u+1][t] + scores[:, s, t]
            if self.beta[s][u] is None:
              self.beta[s][u] = tmp
            else:
              self.beta[s][u] = self.logadd(self.beta[s][u], tmp)
          if u+1 < t+1:
            tmp =  self.beta[s][t] + self.alpha[s][u]  + scores[:, s, t]
            if self.beta[u+1][t] is None:
              self.beta[u+1][t] = tmp
            else:
              self.beta[u+1][t] = self.logadd(self.beta[u+1][t], tmp)

  def _marginal(self, scores):
    batch_size = scores.size(0)
    n = scores.size(1)
    self.log_marginal = [[None for _ in range(n)] for _ in range(n)]
    log_Z = self.alpha[0][n-1]
    for s in range(n):
      for t in np.arange(s, n):
        self.log_marginal[s][t] = self.alpha[s][t] + self.beta[s][t] - log_Z
  
  def _entropy(self, scores):
    batch_size = scores.size(0)
    n = scores.size(1)
    self.entropy = [[None for _ in range(n)] for _ in range(n)]
    for i in range(n):
      self.entropy[i][i] = scores.new(batch_size).fill_(0)
    for k in np.arange(1, n+1):
      for s in range(n):
        t = s + k
        if t > n-1:
          break
        score = []
        prev_ent = []
        for u in np.arange(s, t):
          score.append(self.alpha[s][u] + self.alpha[u+1][t])
          prev_ent.append(self.entropy[s][u] + self.entropy[u+1][t])
        score = torch.stack(score, 1) 
        prev_ent = torch.stack(prev_ent, 1)
        log_prob = F.log_softmax(score, dim = 1)
        prob = log_prob.exp()        
        entropy = ((prev_ent - log_prob)*prob).sum(1)
        self.entropy[s][t] = entropy
      
        
  def _sample(self, scores, alpha = None, argmax = False):    
    # sample from p(tree | sent)
    # also get the spans
    if alpha is None:
      self._forward(scores)
      alpha = self.alpha
    batch_size = scores.size(0)
    n = scores.size(1)
    tree = scores.new(batch_size, n, n).zero_()
    all_log_probs = []
    tree_brackets = []
    spans = []
    for b in range(batch_size):
      sampled = [(0, n-1)]
      span = [(0, n-1)]
      queue = [(0, n-1)] #start, end
      log_probs = []
      tree_str = get_span_str(0, n-1)
      while len(queue) > 0:
        node = queue.pop(0)
        start, end = node
        left_parent = get_span_str(start, None)
        right_parent = get_span_str(None, end)
        score = []
        score_idx = []
        for u in np.arange(start, end):
          score.append(alpha[start][u][b] + alpha[u+1][end][b])
          score_idx.append([(start, u), (u+1, end)])
        score = torch.stack(score, 0) 
        log_prob = F.log_softmax(score, dim = 0)
        if argmax:
          sample = torch.max(log_prob, 0)[1]
        else:
          prob = log_prob.exp()
          sample = torch.multinomial(log_prob.exp(), 1)          
        sample_idx = score_idx[sample.item()]
        log_probs.append(log_prob[sample.item()])
        for idx in sample_idx:
          if idx[0] != idx[1]:
            queue.append(idx)
            span.append(idx)
          sampled.append(idx)
        left_child = '(' + get_span_str(sample_idx[0][0], sample_idx[0][1])    
        right_child = get_span_str(sample_idx[1][0], sample_idx[1][1]) + ')'
        if sample_idx[0][0] != sample_idx[0][1]:
          tree_str = tree_str.replace(left_parent, left_child)
        if sample_idx[1][0] != sample_idx[1][1]:
          tree_str = tree_str.replace(right_parent, right_child)
      all_log_probs.append(torch.stack(log_probs, 0).sum(0))
      tree_brackets.append(tree_str)
      spans.append(span[::-1])
      for idx in sampled:
        tree[b][idx[0]][idx[1]] = 1
        
    all_log_probs = torch.stack(all_log_probs, 0)
    return tree, all_log_probs, tree_brackets, spans

  def _viterbi(self, scores):
    # cky algorithm
    batch_size = scores.size(0)
    n = scores.size(1)
    self.max_scores = scores.new(batch_size, n, n).fill_(-self.huge)
    self.bp = scores.new(batch_size, n, n).zero_()
    self.argmax = scores.new(batch_size, n, n).zero_()
    self.spans = [[] for _ in range(batch_size)]
    tmp = scores.new(batch_size, n).zero_()
    for i in range(n):
      self.max_scores[:, i, i] = scores[:, i, i]      
    for k in np.arange(1, n):
      for s in np.arange(n):
        t = s + k
        if t > n-1:
          break
        for u in np.arange(s, t):
          tmp = self.max_scores[:, s, u] + self.max_scores[:, u+1, t] + scores[:, s, t]
          self.bp[:, s, t][self.max_scores[:, s, t] < tmp] = int(u)
          self.max_scores[:, s, t] = torch.max(self.max_scores[:, s, t], tmp)
    for b in range(batch_size):
      self._backtrack(b, 0, n-1)      
    return self.max_scores[:, 0, n-1], self.argmax, self.spans

  def _backtrack(self, b, s, t):
    u = int(self.bp[b][s][t])
    self.argmax[b][s][t] = 1
    if s == t:
      return None      
    else:
      self.spans[b].insert(0, (s,t))
      self._backtrack(b, s, u)
      self._backtrack(b, u+1, t)
    return None  
 
def get_span_str(start = None, end = None):
  assert(start is not None or end is not None)
  if start is None:
    return ' '  + str(end) + ')'
  elif end is None:
    return '(' + str(start) + ' '
  else:
    return ' (' + str(start) + ' ' + str(end) + ') '