-
Notifications
You must be signed in to change notification settings - Fork 1
/
modules.py
25 lines (22 loc) · 857 Bytes
/
modules.py
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
import torch, torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
class Tsum(nn.Module):
def __init__(self, vocab_size, embed_size):
super().__init__()
self.dec = nn.Sequential(*[
nn.Embedding(vocab_size, embed_size),
nn.LSTM(embed_size, embed_size, bidirectional=True) # think about how choosing hiddensize
])
self.enc = nn.LSTM(embed_size, embed_size)
self.out = nn.Linear(embed_size, vocab_size)
def forward(self, x, y):
embeddings = self.dec(x)
context = torch.zeros_like(embeddings[0])
output = []
while True:
weights = F.cosine_similarity(embeddings, context, dim=1)
context, y = self.enc(context, weights)
output.append(y)
return torch.cat(output, dim=0)