Spaces:
Running
Running
File size: 1,512 Bytes
74d4655 59c6d5c 74d4655 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
import torch
import torch.nn.functional as F
import torch.nn as nn
from .custom_types import Method
class LuongAttention(nn.Module):
def __init__(self, method: Method, hidden_size: int):
super().__init__()
self.hidden_size = hidden_size
if not isinstance(method, Method):
raise ValueError(method, f"should be a member of `Method` enum")
match method:
case Method.DOT:
self.method = self.dot
case Method.GENERAL:
self.method = self.general
self.Wa = nn.Linear(hidden_size, hidden_size)
case Method.CONCAT:
self.method = self.concat
self.Wa = nn.Linear(hidden_size * 2, hidden_size)
self.Va = nn.Parameter(torch.FloatTensor(1, hidden_size))
def dot(self, hidden, encoder_outputs):
return torch.sum(hidden * encoder_outputs, dim=2)
def general(self, hidden, encoder_outputs):
return torch.sum(hidden * self.Wa(encoder_outputs), dim=2)
def concat(self, hidden, encoder_outputs):
hidden = hidden.permute(1, 0, 2)
energy = self.Wa(torch.cat((hidden.permute(1, 0, 2).expand(-1, encoder_outputs.size(1), -1), encoder_outputs), 2)).tanh()
return torch.sum(self.Va * energy, dim=2)
def forward(self, hidden, encoder_outputs):
attn_weights = self.method(hidden, encoder_outputs)
return F.softmax(attn_weights, dim=1).unsqueeze(1) |