| 150 | |
| 151 | |
| 152 | class BiLSTM_CRF(nn.Module): |
| 153 | |
| 154 | def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim): |
| 155 | super(BiLSTM_CRF, self).__init__() |
| 156 | self.embedding_dim = embedding_dim |
| 157 | self.hidden_dim = hidden_dim |
| 158 | self.vocab_size = vocab_size |
| 159 | self.tag_to_ix = tag_to_ix |
| 160 | self.tagset_size = len(tag_to_ix) |
| 161 | |
| 162 | self.word_embeds = nn.Embedding(vocab_size, embedding_dim) |
| 163 | self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2, |
| 164 | num_layers=1, bidirectional=True) |
| 165 | |
| 166 | # Maps the output of the LSTM into tag space. |
| 167 | self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size) |
| 168 | |
| 169 | # Matrix of transition parameters. Entry i,j is the score of |
| 170 | # transitioning *to* i *from* j. |
| 171 | self.transitions = nn.Parameter( |
| 172 | torch.randn(self.tagset_size, self.tagset_size)) |
| 173 | |
| 174 | # These two statements enforce the constraint that we never transfer |
| 175 | # to the start tag and we never transfer from the stop tag |
| 176 | self.transitions.data[tag_to_ix[START_TAG], :] = -10000 |
| 177 | self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000 |
| 178 | |
| 179 | self.hidden = self.init_hidden() |
| 180 | |
| 181 | def init_hidden(self): |
| 182 | return (torch.randn(2, 1, self.hidden_dim // 2), |
| 183 | torch.randn(2, 1, self.hidden_dim // 2)) |
| 184 | |
| 185 | def _forward_alg(self, feats): |
| 186 | # Do the forward algorithm to compute the partition function |
| 187 | init_alphas = torch.full((1, self.tagset_size), -10000.) |
| 188 | # START_TAG has all of the score. |
| 189 | init_alphas[0][self.tag_to_ix[START_TAG]] = 0. |
| 190 | |
| 191 | # Wrap in a variable so that we will get automatic backprop |
| 192 | forward_var = init_alphas |
| 193 | |
| 194 | # Iterate through the sentence |
| 195 | for feat in feats: |
| 196 | alphas_t = [] # The forward tensors at this timestep |
| 197 | for next_tag in range(self.tagset_size): |
| 198 | # broadcast the emission score: it is the same regardless of |
| 199 | # the previous tag |
| 200 | emit_score = feat[next_tag].view( |
| 201 | 1, -1).expand(1, self.tagset_size) |
| 202 | # the ith entry of trans_score is the score of transitioning to |
| 203 | # next_tag from i |
| 204 | trans_score = self.transitions[next_tag].view(1, -1) |
| 205 | # The ith entry of next_tag_var is the value for the |
| 206 | # edge (i -> next_tag) before we do log-sum-exp |
| 207 | next_tag_var = forward_var + trans_score + emit_score |
| 208 | # The forward variable for this tag is log-sum-exp of all the |
| 209 | # scores. |