|
|
import math |
|
|
|
|
|
import einops |
|
|
import torch |
|
|
from torch import nn |
|
|
import torch.nn.functional as F |
|
|
from torch.nn import LayerNorm |
|
|
|
|
|
from common.utils import HiddenData |
|
|
from model.decoder.interaction import BaseInteraction |
|
|
|
|
|
|
|
|
class SlotGatedInteraction(BaseInteraction): |
|
|
def __init__(self, **config): |
|
|
super().__init__(**config) |
|
|
self.intent_linear = nn.Linear(self.config["input_dim"],1, bias=False) |
|
|
self.slot_linear1 = nn.Linear(self.config["input_dim"],1, bias=False) |
|
|
self.slot_linear2 = nn.Linear(self.config["input_dim"],1, bias=False) |
|
|
self.remove_slot_attn = self.config["remove_slot_attn"] |
|
|
self.slot_gate = SlotGate(**config) |
|
|
|
|
|
def forward(self, encode_hidden: HiddenData, **kwargs): |
|
|
input_hidden = encode_hidden.get_slot_hidden_state() |
|
|
|
|
|
seq_lens = encode_hidden.inputs.attention_mask.sum(-1) |
|
|
output_list = [] |
|
|
for index, slen in enumerate(seq_lens): |
|
|
output_list.append(input_hidden[index, slen - 1, :].unsqueeze(0)) |
|
|
intent_input = torch.cat(output_list, dim=0) |
|
|
e_I = torch.tanh(self.intent_linear(intent_input)).squeeze(1) |
|
|
alpha_I = einops.repeat(e_I, 'b -> b h', h=intent_input.shape[-1]) |
|
|
c_I = alpha_I * intent_input |
|
|
intent_hidden = intent_input+c_I |
|
|
if not self.remove_slot_attn: |
|
|
|
|
|
h_k = einops.repeat(self.slot_linear1(input_hidden), 'b l h -> b l c h', c=input_hidden.shape[1]) |
|
|
h_i = einops.repeat(self.slot_linear2(input_hidden), 'b l h -> b l c h', c=input_hidden.shape[1]).transpose(1,2) |
|
|
e_S = torch.tanh(h_k + h_i) |
|
|
alpha_S = torch.softmax(e_S, dim=2).squeeze(3) |
|
|
alpha_S = einops.repeat(alpha_S, 'b l1 l2 -> b l1 l2 h', h=input_hidden.shape[-1]) |
|
|
map_input_hidden = einops.repeat(input_hidden, 'b l h -> b l c h', c=input_hidden.shape[1]) |
|
|
c_S = torch.sum(alpha_S * map_input_hidden, dim=2) |
|
|
else: |
|
|
c_S = input_hidden |
|
|
slot_hidden = input_hidden + c_S * self.slot_gate(c_S,c_I) |
|
|
encode_hidden.update_intent_hidden_state(intent_hidden) |
|
|
encode_hidden.update_slot_hidden_state(slot_hidden) |
|
|
return encode_hidden |
|
|
|
|
|
class SlotGate(nn.Module): |
|
|
def __init__(self, **config): |
|
|
super().__init__() |
|
|
self.linear = nn.Linear(config["input_dim"], config["output_dim"],bias=False) |
|
|
self.v = nn.Parameter(torch.rand(size=[1])) |
|
|
|
|
|
def forward(self, slot_context, intent_context): |
|
|
intent_gate = self.linear(intent_context) |
|
|
intent_gate = einops.repeat(intent_gate, 'b h -> b l h', l=slot_context.shape[1]) |
|
|
return self.v * torch.tanh(slot_context + intent_gate) |
|
|
|
