| 1121 | # |
| 1122 | |
| 1123 | class GreedySearchDecoder(nn.Module): |
| 1124 | def __init__(self, encoder, decoder): |
| 1125 | super(GreedySearchDecoder, self).__init__() |
| 1126 | self.encoder = encoder |
| 1127 | self.decoder = decoder |
| 1128 | |
| 1129 | def forward(self, input_seq, input_length, max_length): |
| 1130 | # Forward input through encoder model |
| 1131 | encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length) |
| 1132 | # Prepare encoder's final hidden layer to be first hidden input to the decoder |
| 1133 | decoder_hidden = encoder_hidden[:self.decoder.n_layers] |
| 1134 | # Initialize decoder input with SOS_token |
| 1135 | decoder_input = torch.ones(1, 1, device=device, dtype=torch.long) * SOS_token |
| 1136 | # Initialize tensors to append decoded words to |
| 1137 | all_tokens = torch.zeros([0], device=device, dtype=torch.long) |
| 1138 | all_scores = torch.zeros([0], device=device) |
| 1139 | # Iteratively decode one word token at a time |
| 1140 | for _ in range(max_length): |
| 1141 | # Forward pass through decoder |
| 1142 | decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs) |
| 1143 | # Obtain most likely word token and its softmax score |
| 1144 | decoder_scores, decoder_input = torch.max(decoder_output, dim=1) |
| 1145 | # Record token and score |
| 1146 | all_tokens = torch.cat((all_tokens, decoder_input), dim=0) |
| 1147 | all_scores = torch.cat((all_scores, decoder_scores), dim=0) |
| 1148 | # Prepare current token to be next decoder input (add a dimension) |
| 1149 | decoder_input = torch.unsqueeze(decoder_input, 0) |
| 1150 | # Return collections of word tokens and scores |
| 1151 | return all_tokens, all_scores |
| 1152 | |
| 1153 | |
| 1154 | ###################################################################### |