
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

torch.manual_seed(1)

class LSTMTagger(nn.Module):

    def __init__(self, embedding_dim, hidden_dim, output_size):
        super(LSTMTagger, self).__init__()
        self.hidden_dim = hidden_dim

        # self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)

        # The LSTM takes word embeddings as inputs, and outputs hidden states
        # with dimensionality hidden_dim.
        self.lstm = nn.LSTM(embedding_dim, hidden_dim)

        # The linear layer that maps from hidden state space to tag space
        self.output = nn.Linear(hidden_dim, output_size)

    def forward(self, input_tensor):
        # embeds = self.word_embeddings(sentence)        
        lstm_out, _ = self.lstm(input_tensor)
        logits = self.output(lstm_out)
        # tag_scores = F.log_softmax(tag_space, dim=1)
        return logits