import math from typing import List import torch.nn as nn import torch.nn.functional as F import torch from torch.nn import init from copy import deepcopy import numpy as np ACTIVATION_POOL = ['relu', 'gelu'] def _get_activation_fn(activation): if activation in ACTIVATION_POOL: return getattr(F, activation) raise RuntimeError("activation should be relu/gelu, not {}".format(activation)) class TreeEncoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward, dropout, max_role_count, activation='gelu', batch_first=False, val_position=False): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first) self._val_position = val_position self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.InstanceNorm1d(d_model) self.norm2 = nn.InstanceNorm1d(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.position_embedding = nn.Embedding(max_role_count, d_model) self.activation = _get_activation_fn(activation) def forward(self, src, src_mask=None, pos_ids=None): """ :param src: concatenation of task embeddings and representation for left and right. src shape: (task_embeddings + left + right, batch_size, dim) :param src_mask: :param pos_ids: :return: """ if len(pos_ids.shape) == 1: sz = src.shape[1] # sz: batch_size pos_ids = pos_ids.unsqueeze(1).expand(-1, sz) # (3, batch_size) position_embedding = self.position_embedding(pos_ids) src2 = self.self_attn(src + position_embedding, src + position_embedding, src, attn_mask=src_mask)[0] src = src + self.dropout1(src2) src = self.norm1(src) src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) src = src + self.dropout2(src2) src = self.norm2(src) return src class InsideEncoder(nn.Module): def __init__(self, config): super().__init__() self.const_size = config.const_size self.left_linear = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, config.const_size)) self.right_linear = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, config.const_size)) # self.const_linear = nn.Sequential(GroupLinear(2, config.hidden_size, config.hidden_size), # nn.GELU(), # GroupLinear(2, config.hidden_size, config.hidden_size)) layer = TreeEncoderLayer(config.hidden_size, config.num_attention_heads, config.intermediate_size, max_role_count=config.max_role_embeddings, dropout=config.attention_probs_dropout_prob, activation='gelu', batch_first=True, val_position=True) self.norm = nn.InstanceNorm1d(config.hidden_size) self.layers = nn.ModuleList([layer] + [deepcopy(layer) for _ in range(config.encoder_num_hidden_layers - 1)]) self._device = None self._pos_ids = None @property def device(self): if self._device is None: self._device = next(self.parameters()).device return self._device @property def pos_ids(self): if self._pos_ids is None: self._pos_ids = torch.arange(2).to(self.device) return self._pos_ids def forward(self, src, span_embeds=None): """ :param src: [batch_size, comb_size, 2, dim] :param span_embeds: [batch_size, dim] :return: """ dim = src.shape[-1] org_shape = src.shape # (batch_size, comb_size, 2, dim) output = src.view(-1, 2, dim) # torch.cuda.synchronize() # with torch.cuda.stream(self.s1): left_const = self.left_linear(output[:, 0, :]) right_const = self.right_linear(output[:, 1, :]) for mod in self.layers: output = mod(output, pos_ids=self.pos_ids.unsqueeze(0)) mat_scores = torch.einsum("bi,bi->b", left_const, right_const) / math.sqrt(self.const_size) mat_scores = mat_scores.view(*org_shape[:-2]) # (batch_size, comb_size) if span_embeds is not None: output = output.sum(dim=1).view(*org_shape[:-2], dim) + span_embeds.unsqueeze(1) else: output = output.sum(dim=1).view(*org_shape[:-2], dim) return mat_scores, self.norm(output) class OutsideEncoder(nn.Module): def __init__(self, config) -> None: super().__init__() self.parent_linear = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, config.const_size)) self.left_linear = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, config.const_size)) self.right_linear = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, config.const_size)) # self.const_linear = nn.Sequential(GroupLinear(3, config.hidden_size, config.hidden_size), # nn.GELU(), # GroupLinear(3, config.hidden_size, config.hidden_size)) self.const_size = config.const_size layer = TreeEncoderLayer(config.hidden_size, config.num_attention_heads, config.intermediate_size, max_role_count=config.max_role_embeddings, dropout=config.attention_probs_dropout_prob, activation='gelu', batch_first=True, val_position=True) self.norm = nn.InstanceNorm1d(config.hidden_size) self.layers = nn.ModuleList([layer] + [deepcopy(layer) for _ in range(config.decoder_num_hidden_layers - 1)]) self._device = None self._dec_pos_ids = None @property def device(self): if self._device is None: self._device = next(self.parameters()).device return self._device @property def dec_pos_ids(self): if self._dec_pos_ids is None: self._dec_pos_ids = torch.tensor([[0, 1], [0, 2]], device=self.device) return self._dec_pos_ids def forward(self, parent_ij, child_ikj, parent_scores=None, child_scores=None): """ :param parent_ij: (batch_size, dim) :param parent_scores: (batch_size, 1) :param child_ikj: (batch_size, comb_size, 2, dim) :param child_scores: (batch_size, comb_size, 2) :return: (batch_size, 2), (batch_size, 2, dim) """ # p_l = parent_ij @ self.W_outside_r # (batch_size, dim) # out_score_ik = torch.einsum('bd, bcd->bc', p_l, child_ikj[:, :, 1, :]) # (batch_size, comb_size) # out_score_ik = out_score_ik + parent_scores + child_scores[:, :, 1, 0] # (batch_size, comb_size) # p_r = parent_ij @ self.W_outside_l # (batch_size, dim) # out_score_kj = torch.einsum('bd, bcd->bc', p_r, child_ikj[:, :, 0, :]) # (batch_size, comb_size) # out_score_kj = out_score_kj + parent_scores + child_scores[:, :, 0, 0] # (batch_size, comb_size) batch_size = child_ikj.shape[0] comb_size = child_ikj.shape[1] comb_size = child_ikj.shape[1] parent_ij_ext = parent_ij.unsqueeze(1).unsqueeze(2).repeat(1, comb_size, 2, 1) # (batch_size, comb_size, 2, dim) inputs = torch.stack([parent_ij_ext, child_ikj.flip([2])], dim=3) # (batch_size, comb_size, 2, 2, dim) inputs = inputs.view(batch_size * comb_size * 2, 2, -1) outside_scores = None if parent_scores is not None and child_scores is not None: parent_const = self.parent_linear(inputs[:, 0, :]) right_child_const = self.right_linear(inputs[::2, 1, :]) left_child_const = self.left_linear(inputs[1::2, 1, :]) # self.dec_pos_ids: (2, 2) pos_ids = self.dec_pos_ids.repeat(batch_size * comb_size, 1) for mod in self.layers: inputs = mod(inputs, pos_ids=pos_ids) if parent_scores is not None and child_scores is not None: parent_const_r = parent_const[::2, :] parent_const_l = parent_const[1::2, :] left_score = (parent_const_r * right_child_const).sum(dim=-1) / math.sqrt(self.const_size) # right_score = torch.einsum('bi,bi->b', parent_const_l, left_child_const) / math.sqrt(self.const_size) right_score = (parent_const_l * left_child_const).sum(dim=-1) / math.sqrt(self.const_size) left_score = left_score.view(batch_size, comb_size) right_score = right_score.view(batch_size, comb_size) out_score_ik = left_score out_score_kj = right_score outside_scores = torch.stack([out_score_ik, out_score_kj], dim=2) # inputs: (?, 2, dim) out_e_ij = self.norm(inputs.sum(dim=1)) return outside_scores, \ out_e_ij.view(batch_size, comb_size, 2, -1) # (batch_size, comb_size, 2, dim)