Update structformer_as_hf.py

structformer_as_hf.py

@@ -6,6 +6,13 @@ from transformers import PreTrainedModel
 from transformers import PretrainedConfig
 from transformers.modeling_outputs import MaskedLMOutput
 from typing import List
+from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    MaskedLMOutput,
+    SequenceClassifierOutput
+)
 
 ##########################################
 # HuggingFace Config
@@ -67,7 +74,6 @@ class Conv1d(nn.Module):
 
   def __init__(self, hidden_size, kernel_size, dilation=1):
     """Initialization.
-
     Args:
       hidden_size: dimension of input embeddings
       kernel_size: convolution kernel size
@@ -90,7 +96,6 @@ class Conv1d(nn.Module):
 
   def forward(self, x):
     """Compute convolution.
-
     Args:
       x: input embeddings
     Returns:
@@ -114,7 +119,6 @@ class MultiheadAttention(nn.Module):
                out_proj=True,
                relative_bias=True):
     """Initialization.
-
     Args:
       embed_dim: dimension of input embeddings
       num_heads: number of self-attention heads
@@ -174,7 +178,6 @@ class MultiheadAttention(nn.Module):
 
   def forward(self, query, key_padding_mask=None, attn_mask=None):
     """Compute multi-head self-attention.
-
     Args:
       query: input embeddings
       key_padding_mask: 3D mask that prevents attention to certain positions
@@ -254,7 +257,6 @@ class TransformerLayer(nn.Module):
                activation="leakyrelu",
                relative_bias=True):
     """Initialization.
-
     Args:
       d_model: dimension of inputs
       nhead: number of self-attention heads
@@ -285,7 +287,6 @@ class TransformerLayer(nn.Module):
 
   def forward(self, src, attn_mask=None, key_padding_mask=None):
     """Pass the input through the encoder layer.
-
     Args:
       src: the sequence to the encoder layer (required).
       attn_mask: the mask for the src sequence (optional).
@@ -301,6 +302,30 @@ class TransformerLayer(nn.Module):
 
     return src3
 
+
+
+class RobertaClassificationHead(nn.Module):
+    """Head for sentence-level classification tasks."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        classifier_dropout = (
+            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+        )
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
+
+    def forward(self, features, **kwargs):
+        x = features[:, 0, :]  # take <s> token (equiv. to [CLS])
+        x = self.dropout(x)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x
+
+
 ##########################################
 # Custom Models
 ##########################################
@@ -362,7 +387,6 @@ class Transformer(nn.Module):
                pos_emb=False,
                pad=0):
     """Initialization.
-
     Args:
       hidden_size: dimension of inputs and hidden states
       nlayers: number of layers
@@ -437,7 +461,6 @@ class Transformer(nn.Module):
 
   def forward(self, x, pos):
     """Pass the input through the encoder layer.
-
     Args:
       x: input tokens (required).
       pos: position for each token (optional).
@@ -474,7 +497,6 @@ class StructFormer(Transformer):
                relations=('head', 'child'),
                weight_act='softmax'):
     """Initialization.
-
     Args:
       hidden_size: dimension of inputs and hidden states
       nlayers: number of layers
@@ -548,7 +570,6 @@ class StructFormer(Transformer):
 
   def parse(self, x, pos, embeds=None):
     """Parse input sentence.
-
     Args:
       x: input tokens (required).
       pos: position for each token (optional).
@@ -735,6 +756,300 @@ class StructFormer(Transformer):
             attentions=None,
     )   
     
+
+
+
+class StructFormerClassification(Transformer):
+  """StructFormer model."""
+
+  def __init__(self,
+               hidden_size,
+               n_context_layers,
+               nlayers,
+               ntokens,
+               nhead=8,
+               dropout=0.1,
+               dropatt=0.1,
+               relative_bias=False,
+               pos_emb=False,
+               pad=0,
+               n_parser_layers=4,
+               conv_size=9,
+               relations=('head', 'child'),
+               weight_act='softmax',
+               config=None,
+               ):
+    
+
+    super(StructFormerClassification, self).__init__(
+        hidden_size,
+        nlayers,
+        ntokens,
+        nhead=nhead,
+        dropout=dropout,
+        dropatt=dropatt,
+        relative_bias=relative_bias,
+        pos_emb=pos_emb,
+        pad=pad)
+    
+    self.num_labels = config.num_labels
+    self.config = config
+
+    if n_context_layers > 0:
+      self.context_layers = nn.ModuleList([
+          TransformerLayer(hidden_size, nhead, hidden_size * 4, dropout,
+                                  dropatt=dropatt, relative_bias=relative_bias)
+          for _ in range(n_context_layers)])
+  
+    self.parser_layers = nn.ModuleList([
+        nn.Sequential(Conv1d(hidden_size, conv_size),
+                      nn.LayerNorm(hidden_size, elementwise_affine=False),
+                      nn.Tanh()) for i in range(n_parser_layers)])
+
+    self.distance_ff = nn.Sequential(
+        Conv1d(hidden_size, 2),
+        nn.LayerNorm(hidden_size, elementwise_affine=False), nn.Tanh(),
+        nn.Linear(hidden_size, 1))
+
+    self.height_ff = nn.Sequential(
+        nn.Linear(hidden_size, hidden_size),
+        nn.LayerNorm(hidden_size, elementwise_affine=False), nn.Tanh(),
+        nn.Linear(hidden_size, 1))
+
+    n_rel = len(relations)
+    self._rel_weight = nn.Parameter(torch.zeros((nlayers, nhead, n_rel)))
+    self._rel_weight.data.normal_(0, 0.1)
+
+    self._scaler = nn.Parameter(torch.zeros(2))
+
+    self.n_parse_layers = n_parser_layers
+    self.n_context_layers = n_context_layers
+    self.weight_act = weight_act
+    self.relations = relations
+    
+    self.classifier = RobertaClassificationHead(config)
+
+  @property
+  def scaler(self):
+    return self._scaler.exp()
+
+  @property
+  def rel_weight(self):
+    if self.weight_act == 'sigmoid':
+      return torch.sigmoid(self._rel_weight)
+    elif self.weight_act == 'softmax':
+      return torch.softmax(self._rel_weight, dim=-1)
+
+  def parse(self, x, pos, embeds=None):
+    """Parse input sentence.
+    Args:
+      x: input tokens (required).
+      pos: position for each token (optional).
+    Returns:
+      distance: syntactic distance
+      height: syntactic height
+    """
+
+    mask = (x != self.pad)
+    mask_shifted = F.pad(mask[:, 1:], (0, 1), value=0)
+
+    
+    if embeds is not None:
+      h = embeds
+    else:
+      h = self.emb(x)
+    
+    for i in range(self.n_parse_layers):
+      h = h.masked_fill(~mask[:, :, None], 0)
+      h = self.parser_layers[i](h)
+
+    height = self.height_ff(h).squeeze(-1)
+    height.masked_fill_(~mask, -1e9)
+
+    distance = self.distance_ff(h).squeeze(-1)
+    distance.masked_fill_(~mask_shifted, 1e9)
+
+    # Calbrating the distance and height to the same level
+    length = distance.size(1)
+    height_max = height[:, None, :].expand(-1, length, -1)
+    height_max = torch.cummax(
+        height_max.triu(0) - torch.ones_like(height_max).tril(-1) * 1e9,
+        dim=-1)[0].triu(0)
+
+    margin_left = torch.relu(
+        F.pad(distance[:, :-1, None], (0, 0, 1, 0), value=1e9) - height_max)
+    margin_right = torch.relu(distance[:, None, :] - height_max)
+    margin = torch.where(margin_left > margin_right, margin_right,
+                         margin_left).triu(0)
+
+    margin_mask = torch.stack([mask_shifted] + [mask] * (length - 1), dim=1)
+    margin.masked_fill_(~margin_mask, 0)
+    margin = margin.max()
+
+    distance = distance - margin
+
+    return distance, height
+
+  def compute_block(self, distance, height):
+    """Compute constituents from distance and height."""
+
+    beta_logits = (distance[:, None, :] - height[:, :, None]) * self.scaler[0]
+
+    gamma = torch.sigmoid(-beta_logits)
+    ones = torch.ones_like(gamma)
+
+    block_mask_left = cummin(
+        gamma.tril(-1) + ones.triu(0), reverse=True, max_value=1)
+    block_mask_left = block_mask_left - F.pad(
+        block_mask_left[:, :, :-1], (1, 0), value=0)
+    block_mask_left.tril_(0)
+
+    block_mask_right = cummin(
+        gamma.triu(0) + ones.tril(-1), exclusive=True, max_value=1)
+    block_mask_right = block_mask_right - F.pad(
+        block_mask_right[:, :, 1:], (0, 1), value=0)
+    block_mask_right.triu_(0)
+
+    block_p = block_mask_left[:, :, :, None] * block_mask_right[:, :, None, :]
+    block = cumsum(block_mask_left).tril(0) + cumsum(
+        block_mask_right, reverse=True).triu(1)
+
+    return block_p, block
+
+  def compute_head(self, height):
+    """Estimate head for each constituent."""
+
+    _, length = height.size()
+    head_logits = height * self.scaler[1]
+    index = torch.arange(length, device=height.device)
+
+    mask = (index[:, None, None] <= index[None, None, :]) * (
+        index[None, None, :] <= index[None, :, None])
+    head_logits = head_logits[:, None, None, :].repeat(1, length, length, 1)
+    head_logits.masked_fill_(~mask[None, :, :, :], -1e9)
+
+    head_p = torch.softmax(head_logits, dim=-1)
+
+    return head_p
+
+  def generate_mask(self, x, distance, height):
+    """Compute head and cibling distribution for each token."""
+
+    bsz, length = x.size()
+
+    eye = torch.eye(length, device=x.device, dtype=torch.bool)
+    eye = eye[None, :, :].expand((bsz, -1, -1))
+
+    block_p, block = self.compute_block(distance, height)
+    head_p = self.compute_head(height)
+    head = torch.einsum('blij,bijh->blh', block_p, head_p)
+    head = head.masked_fill(eye, 0)
+    child = head.transpose(1, 2)
+    cibling = torch.bmm(head, child).masked_fill(eye, 0)
+
+    rel_list = []
+    if 'head' in self.relations:
+      rel_list.append(head)
+    if 'child' in self.relations:
+      rel_list.append(child)
+    if 'cibling' in self.relations:
+      rel_list.append(cibling)
+
+    rel = torch.stack(rel_list, dim=1)
+
+    rel_weight = self.rel_weight
+
+    dep = torch.einsum('lhr,brij->lbhij', rel_weight, rel)
+    att_mask = dep.reshape(self.nlayers, bsz * self.nhead, length, length)
+
+    return att_mask, cibling, head, block
+
+  def encode(self, x, pos, att_mask=None, context_layers=False):
+    """Structformer encoding process."""
+
+    if context_layers:
+      """Standard transformer encode process."""
+      h = self.emb(x)
+      if hasattr(self, 'pos_emb'):
+        h = h + self.pos_emb(pos)
+      h_list = []
+      visibility = self.visibility(x, x.device)
+      for i in range(self.n_context_layers):
+        h_list.append(h)
+        h = self.context_layers[i](
+            h.transpose(0, 1), key_padding_mask=visibility).transpose(0, 1)
+
+      output = h
+      h_array = torch.stack(h_list, dim=2)
+      return output
+    
+    else:
+      visibility = self.visibility(x, x.device)
+      h = self.emb(x)
+      if hasattr(self, 'pos_emb'):
+        assert pos.max() < 500
+        h = h + self.pos_emb(pos)
+      for i in range(self.nlayers):
+        h = self.layers[i](
+            h.transpose(0, 1), attn_mask=att_mask[i],
+            key_padding_mask=visibility).transpose(0, 1)
+      return h
+
+  def forward(self, input_ids, labels=None, position_ids=None, **kwargs):
+        
+    x = input_ids
+    batch_size, length = x.size()
+
+    if position_ids is None:
+      pos = torch.arange(length, device=x.device).expand(batch_size, length)
+
+    context_layers_output = None
+    if self.n_context_layers > 0:
+      context_layers_output = self.encode(x, pos, context_layers=True)    
+    
+    distance, height = self.parse(x, pos, embeds=context_layers_output)
+    att_mask, cibling, head, block = self.generate_mask(x, distance, height)
+    
+    raw_output = self.encode(x, pos, att_mask)
+    raw_output = self.norm(raw_output)
+    raw_output = self.drop(raw_output)
+
+    #output = self.output_layer(raw_output)
+    logits = self.classifier(raw_output)
+    
+    loss = None
+    if labels is not None:
+        if self.config.problem_type is None:
+            if self.num_labels == 1:
+                self.config.problem_type = "regression"
+            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                self.config.problem_type = "single_label_classification"
+            else:
+                self.config.problem_type = "multi_label_classification"
+
+        if self.config.problem_type == "regression":
+            loss_fct = MSELoss()
+            if self.num_labels == 1:
+                loss = loss_fct(logits.squeeze(), labels.squeeze())
+            else:
+                loss = loss_fct(logits, labels)
+        elif self.config.problem_type == "single_label_classification":
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+        elif self.config.problem_type == "multi_label_classification":
+            loss_fct = BCEWithLogitsLoss()
+            loss = loss_fct(logits, labels)
+
+
+    return SequenceClassifierOutput(
+        loss=loss,
+        logits=logits,
+        hidden_states=None,
+        attentions=None,
+    )
+
+
+
 ##########################################
 # HuggingFace Model
 ##########################################
@@ -760,5 +1075,49 @@ class StructformerModel(PreTrainedModel):
                weight_act=config.weight_act
         )
 
+    def forward(self, input_ids, labels=None, **kwargs):
+        return self.model(input_ids, labels=labels, **kwargs)
+
+
+
+class StructformerModelForSequenceClassification(PreTrainedModel):
+    config_class = StructformerConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = StructFormerClassification(
+               hidden_size=config.hidden_size,
+               n_context_layers=config.n_context_layers,
+               nlayers=config.nlayers,
+               ntokens=config.ntokens,
+               nhead=config.nhead,
+               dropout=config.dropout,
+               dropatt=config.dropatt,
+               relative_bias=config.relative_bias,
+               pos_emb=config.pos_emb,
+               pad=config.pad,
+               n_parser_layers=config.n_parser_layers,
+               conv_size=config.conv_size,
+               relations=config.relations,
+               weight_act=config.weight_act,
+               config=config)
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            if module.bias is not None:
+                module.bias.data.zero_()
+                module.weight.data.fill_(1.0)
+
+    
     def forward(self, input_ids, labels=None, **kwargs):
         return self.model(input_ids, labels=labels, **kwargs)