| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class SGNSModel(nn.Module): |
| def __init__(self, vocab_size, embedding_dim): |
| super(SGNSModel, self).__init__() |
| self.target_embedding = nn.Embedding(vocab_size, embedding_dim) |
| self.context_embedding = nn.Embedding(vocab_size, embedding_dim) |
|
|
| def forward(self, target, context): |
| target_embed = self.target_embedding(target) |
| context_embed = self.context_embedding(context) |
| dot = torch.einsum('ij,ji->i',target_embed,context_embed.t()) |
| output_layer = F.sigmoid(dot) |
| return output_layer |
| |
|
|
| class NWPModel(nn.Module): |
| def __init__(self, vocab_size, embedding_dim, context_size, hidden_size, sgns_model): |
| super(NWPModel, self).__init__() |
| self.embedding = nn.Embedding(vocab_size, embedding_dim) |
| self.embedding.weight = nn.Parameter(sgns_model.target_embedding.weight.data) |
| self.embedding.weight.requires_grad = False |
| self.hidden = nn.Linear(embedding_dim*context_size, hidden_size) |
| self.activation = nn.ReLU() |
| self.output = nn.Linear(hidden_size, vocab_size) |
|
|
| def forward(self, context): |
| embed = self.embedding(context) |
| if context.ndim != 1: |
| embed = torch.flatten(embed, start_dim=1) |
| else: |
| embed = torch.flatten(embed) |
| hidden = self.hidden(embed) |
| activation = self.activation(hidden) |
| output = self.output(activation) |
| return output |
|
|
| class NWPModelScratch(nn.Module): |
| def __init__(self, vocab_size, embedding_dim, context_size, hidden_size): |
| super(NWPModelScratch, self).__init__() |
| self.embedding = nn.Embedding(vocab_size, embedding_dim) |
| self.hidden = nn.Linear(embedding_dim*context_size, hidden_size) |
| self.activation = nn.ReLU() |
| self.output = nn.Linear(hidden_size, vocab_size) |
|
|
| def forward(self, context): |
| embed = self.embedding(context) |
| if context.ndim != 1: |
| embed = torch.flatten(embed, start_dim=1) |
| else: |
| embed = torch.flatten(embed) |
| hidden = self.hidden(embed) |
| activation = self.activation(hidden) |
| output = self.output(activation) |
| return output |
|
|
