#!/usr/bin/env python

import torch
from torch import nn

class MLPModel(nn.Module):
  def __init__(self, vocab_sz, hidden_dim, dropout_p):
    super(MLPModel, self).__init__()
    
    self.fc1 = nn.Linear(in_features=vocab_sz, out_features=hidden_dim)
    self.relu = nn.ReLU()
    self.dropout = nn.Dropout(dropout_p)
    self.fc2 = nn.Linear(in_features=hidden_dim, out_features=1)

  def forward(self, x_in):
    x_out = self.fc1(x_in)
    x_out = self.dropout(self.relu(x_out))
    x_out = self.fc2(x_out)
    
    return x_out.squeeze(1)