Method forward

model/embedding.py:169–211

(self, vocab_ids)

            fan_mode=fan_mode)

    def forward(self, vocab_ids):
        seq_length = vocab_ids.size(1)
        actual_length = vocab_ids.size(1) - self.radius * 2
        trim_vocab_id = vocab_ids[:, self.radius:seq_length - self.radius]
        slice_vocabs = \
            [vocab_ids[:, i:i + self.region_size] for i in
             range(actual_length)]
        slice_vocabs = torch.cat(slice_vocabs, 1)
        slice_vocabs = \
            slice_vocabs.view(-1, actual_length, self.region_size)

        if self.region_embedding_type == RegionEmbeddingType.WC:
            vocab_embedding = self.embedding(slice_vocabs)
            context_embedding = self.context_embedding(trim_vocab_id)
            context_embedding = context_embedding.view(
                -1, actual_length, self.region_size, self.embedding_dim)
            region_embedding = vocab_embedding * context_embedding
            region_embedding, _ = region_embedding.max(2)
        elif self.region_embedding_type == RegionEmbeddingType.CW:
            vocab_embedding = self.embedding(trim_vocab_id).unsqueeze(2)
            context_embedding = self.context_embedding(slice_vocabs)
            size = context_embedding.size()
            context_embedding = context_embedding.view(
                size[0], size[1], size[2], self.region_size, self.embedding_dim)
            mask = torch.ones(
                [self.region_size, self.region_size, self.embedding_dim])

            for i in range(self.region_size):
                mask[i][self.region_size - i - 1] = 0.
            neg_mask = mask * -65500.0
            mask = mask.le(0).float()
            mask = mask.unsqueeze(0).unsqueeze(0)
            context_embedding = context_embedding * mask
            context_embedding = context_embedding + neg_mask
            context_embedding, _ = context_embedding.max(3)
            region_embedding = vocab_embedding * context_embedding
            region_embedding, _ = region_embedding.max(2)
        else:
            raise TypeError(
                "Unsupported region embedding type: %s." %
                self.region_embedding_type)

        return region_embedding

class PositionEmbedding(torch.nn.Module):
    ''' Reference: attention is all you need '''