import torch
import torch.nn as nn

class Seq2Seq(nn.Module):
    def __init__(self, src_vocab, tgt_vocab, embed_dim=256, hidden_dim=512):
        super().__init__()
        self.encoder = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
        self.decoder = nn.LSTM(embed_dim + hidden_dim*2, hidden_dim, batch_first=True)
        self.attention = nn.Linear(hidden_dim*3, hidden_dim)
        self.v = nn.Linear(hidden_dim, 1, bias=False)
        self.out = nn.Linear(hidden_dim*3 + embed_dim, tgt_vocab)
    def attention_forward(self, hidden, encoder_outputs):
        src_len = encoder_outputs.shape[1]
        hidden = hidden.repeat(src_len, 1, 1).permute(1,0,2)
        energy = torch.tanh(self.attention(torch.cat([hidden, encoder_outputs], dim=2)))
        attention = self.v(energy).squeeze(2)
        return torch.softmax(attention, dim=1)