Commit 4afa455e authored by Haoran Li's avatar Haoran Li Committed by Facebook Github Bot
Browse files

make fairseq models compatible with character inputs and use character inputs for elmo in pytext 

Summary: Pull Request resolved: https://github.com/pytorch/fairseq/pull/321

Reviewed By: alexeib

Differential Revision: D10430186

fbshipit-source-id: 9cc8fe0f202cc49370cecf36312bcc9bf0b4deee
parent 613ffeea

fairseq/modules/character_token_embedder.py

+24 −10
Original line number Diff line number Diff line
@@ -17,6 +17,9 @@ from typing import List, Tuple 
from .highway import Highway

from fairseq.data import Dictionary



CHAR_PAD_IDX = 0

CHAR_EOS_IDX = 257





class CharacterTokenEmbedder(torch.nn.Module):

    def __init__(
@@ -27,13 +30,16 @@ class CharacterTokenEmbedder(torch.nn.Module): 
            word_embed_dim: int,

            highway_layers: int,

            max_char_len: int = 50,

            char_inputs: bool = False

    ):

        super(CharacterTokenEmbedder, self).__init__()



        self.embedding_dim = word_embed_dim

        self.max_char_len = max_char_len

        self.char_embeddings = nn.Embedding(257, char_embed_dim, padding_idx=0)

        self.symbol_embeddings = nn.Parameter(torch.FloatTensor(2, word_embed_dim))

        self.eos_idx, self.unk_idx = 0, 1

        self.char_inputs = char_inputs



        self.convolutions = nn.ModuleList()

        for width, out_c in filters:
@@ -84,26 +90,34 @@ class CharacterTokenEmbedder(torch.nn.Module): 


    def forward(

            self,

            words: torch.Tensor,

            input: torch.Tensor,

    ):

        self.word_to_char = self.word_to_char.type_as(words)



        flat_words = words.view(-1)

        word_embs = self._convolve(self.word_to_char[flat_words])



        if self.char_inputs:

            chars = input.view(-1, self.max_char_len)

            pads = chars[:, 0].eq(CHAR_PAD_IDX)

            eos = chars[:, 0].eq(CHAR_EOS_IDX)

            if eos.any():

                chars[eos] = 0

            unk = None

        else:

            self.word_to_char = self.word_to_char.type_as(input)

            flat_words = input.view(-1)

            chars = self.word_to_char[flat_words]

            pads = flat_words.eq(self.vocab.pad())

            eos = flat_words.eq(self.vocab.eos())

            unk = flat_words.eq(self.vocab.unk())



        word_embs = self._convolve(chars)

        if pads.any():

            word_embs[pads] = 0



        eos = flat_words.eq(self.vocab.eos())

        if eos.any():

            word_embs[eos] = self.symbol_embeddings[self.eos_idx]



        unk = flat_words.eq(self.vocab.unk())

        if unk.any():

        if unk is not None and unk.any():

            word_embs[unk] = self.symbol_embeddings[self.unk_idx]



        return word_embs.view(words.size() + (-1,))

        return word_embs.view(input.size()[:2] + (-1,))



    def _convolve(

            self,