SentEmbed/model.py at master · Vansil/SentEmbed · GitHub

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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import torch.nn as nn
import torch


class Combine(nn.Module):
    '''
    Layer that takes two sentence embeddings and combines them
    Concatenates: sentence embeddings, absolute difference, element-wise product
    '''
    def forward(self, premise, hypothesis):
        absdiff = abs(premise-hypothesis)
        product = premise * hypothesis
        out = torch.cat((premise, hypothesis, absdiff, product), dim=1)

        return out


class Average(nn.Module):
    '''
    Layer that averages its input in dimension 1
    Used for the Baseline sentence embedding
    '''
    def forward(self, x):
        return x.mean(dim=1)

class UniLstmNet(nn.Module):
    '''
    Unidirectional LSTM network, the last hidden state is the sentence representation
    '''
    def __init__(self, embed_weights):
        super(UniLstmNet,self).__init__()

        hidden_dim = 300
        self.embed = nn.Embedding.from_pretrained(embed_weights)
        self.lstm = nn.LSTM(self.embed.embedding_dim, hidden_dim)

        self.comb  = Combine()

        self.sequential = nn.Sequential(
            nn.Linear(1200,512),
            nn.ReLU(),
            nn.Linear(512,3)
        )

    def forward(self, premise, hypothesis):
        embed_premise = self.encode(premise)
        embed_hypothesis = self.encode(hypothesis)

        # Classifier
        combined = self.comb(embed_premise, embed_hypothesis)
        out = self.sequential(combined)

        return out

    def encode(self, sentence):
        '''
        Encode a sentence for SentEval evaluation
        '''
        # GloVe embedding
        embed = self.embed(sentence)
        embed = embed.permute([1,0,2])

        # LSTM
        hidden_dim = 300
        batch_size = sentence.shape[0]
        device_name = 'cuda' if torch.cuda.is_available() else 'cpu'
        device = torch.device(device_name)
        hidden_0 = torch.zeros(1, batch_size, hidden_dim).to(device)
        state_0  = torch.zeros(1, batch_size, hidden_dim).to(device)
        _, (hidden, _) = self.lstm(embed, (hidden_0, state_0))

        return hidden.squeeze()

    def trainable_params(self):
        return [p for p in self.lstm.parameters()] + [p for p in self.sequential.parameters()]


class BaselineNet(nn.Module):
    '''
    Baseline network with one hidden layer of 100 neurons
    '''
    def __init__(self, embed_weights):
        super(BaselineNet,self).__init__()

        self.embed = nn.Embedding.from_pretrained(embed_weights)
        self.avg   = Average()
        self.comb  = Combine()

        self.sequential = nn.Sequential(
            nn.Linear(1200,512),
            nn.ReLU(),
            nn.Linear(512,3)
        )

    def forward(self, premise, hypothesis):
        embed_premise    = self.avg(self.embed(premise))
        embed_hypothesis = self.avg(self.embed(hypothesis))

        combined = self.comb(embed_premise, embed_hypothesis)
        out = self.sequential(combined)

        return out

    def encode(self, sentence):
        '''
        Encode a sentence for SentEval evaluation
        '''
        return self.avg(self.embed(sentence))

    def trainable_params(self):
        return [p for p in self.sequential.parameters()]