Add pretrained embedding support (#151)

fairseq/models/fconv.py

@@ -30,10 +30,14 @@ class FConvModel(FairseqModel):
                             help='dropout probability')
         parser.add_argument('--encoder-embed-dim', type=int, metavar='N',
                             help='encoder embedding dimension')
+        parser.add_argument('--encoder-embed-path', default=None, type=str, metavar='STR',
+                            help='path to pre-trained encoder embedding')
         parser.add_argument('--encoder-layers', type=str, metavar='EXPR',
                             help='encoder layers [(dim, kernel_size), ...]')
         parser.add_argument('--decoder-embed-dim', type=int, metavar='N',
                             help='decoder embedding dimension')
+        parser.add_argument('--decoder-embed-path', default=None, type=str, metavar='STR',
+                            help='path to pre-trained decoder embedding')
         parser.add_argument('--decoder-layers', type=str, metavar='EXPR',
                             help='decoder layers [(dim, kernel_size), ...]')
         parser.add_argument('--decoder-out-embed-dim', type=int, metavar='N',
@@ -53,9 +57,21 @@ class FConvModel(FairseqModel):
             args.max_target_positions = args.max_positions
         if not hasattr(args, 'share_input_output_embed'):
             args.share_input_output_embed = False
+
+        encoder_embed_dict = None
+        if args.encoder_embed_path:
+            encoder_embed_dict = utils.parse_embedding(args.encoder_embed_path)
+            utils.print_embed_overlap(encoder_embed_dict, src_dict)
+
+        decoder_embed_dict = None
+        if args.decoder_embed_path:
+            decoder_embed_dict = utils.parse_embedding(args.decoder_embed_path)
+            utils.print_embed_overlap(decoder_embed_dict, dst_dict)
+
         encoder = FConvEncoder(
             src_dict,
             embed_dim=args.encoder_embed_dim,
+            embed_dict=encoder_embed_dict,
             convolutions=eval(args.encoder_layers),
             dropout=args.dropout,
             max_positions=args.max_source_positions,
@@ -63,6 +79,7 @@ class FConvModel(FairseqModel):
         decoder = FConvDecoder(
             dst_dict,
             embed_dim=args.decoder_embed_dim,
+            embed_dict=decoder_embed_dict,
             convolutions=eval(args.decoder_layers),
             out_embed_dim=args.decoder_out_embed_dim,
             attention=eval(args.decoder_attention),
@@ -75,8 +92,8 @@ class FConvModel(FairseqModel):
 
 class FConvEncoder(FairseqEncoder):
     """Convolutional encoder"""
-    def __init__(self, dictionary, embed_dim=512, max_positions=1024,
-                 convolutions=((512, 3),) * 20, dropout=0.1):
+    def __init__(self, dictionary, embed_dim=512, embed_dict=None,
+                 max_positions=1024, convolutions=((512, 3),) * 20, dropout=0.1):
         super().__init__(dictionary)
         self.dropout = dropout
         self.num_attention_layers = None
@@ -84,6 +101,9 @@ class FConvEncoder(FairseqEncoder):
         num_embeddings = len(dictionary)
         padding_idx = dictionary.pad()
         self.embed_tokens = Embedding(num_embeddings, embed_dim, padding_idx)
+        if embed_dict:
+            self.embed_tokens = utils.load_embedding(embed_dict, self.dictionary, self.embed_tokens)
+
         self.embed_positions = PositionalEmbedding(
             max_positions,
             embed_dim,
@@ -197,7 +217,8 @@ class AttentionLayer(nn.Module):
 
 class FConvDecoder(FairseqIncrementalDecoder):
     """Convolutional decoder"""
-    def __init__(self, dictionary, embed_dim=512, out_embed_dim=256,
+    def __init__(self, dictionary, embed_dim=512,
+                 embed_dict=None, out_embed_dim=256,
                  max_positions=1024, convolutions=((512, 3),) * 20,
                  attention=True, dropout=0.1, share_embed=False):
         super().__init__(dictionary)
@@ -215,6 +236,9 @@ class FConvDecoder(FairseqIncrementalDecoder):
         num_embeddings = len(dictionary)
         padding_idx = dictionary.pad()
         self.embed_tokens = Embedding(num_embeddings, embed_dim, padding_idx)
+        if embed_dict:
+            self.embed_tokens = utils.load_embedding(embed_dict, self.dictionary, self.embed_tokens)
+
         self.embed_positions = PositionalEmbedding(
             max_positions,
             embed_dim,

fairseq/models/lstm.py

@@ -28,10 +28,14 @@ class LSTMModel(FairseqModel):
                             help='dropout probability')
         parser.add_argument('--encoder-embed-dim', type=int, metavar='N',
                             help='encoder embedding dimension')
+        parser.add_argument('--encoder-embed-path', default=None, type=str, metavar='STR',
+                            help='path to pre-trained encoder embedding')
         parser.add_argument('--encoder-layers', type=int, metavar='N',
                             help='number of encoder layers')
         parser.add_argument('--decoder-embed-dim', type=int, metavar='N',
                             help='decoder embedding dimension')
+        parser.add_argument('--decoder-embed-path', default=None, type=str, metavar='STR',
+                            help='path to pre-trained decoder embedding')
         parser.add_argument('--decoder-layers', type=int, metavar='N',
                             help='number of decoder layers')
         parser.add_argument('--decoder-out-embed-dim', type=int, metavar='N',
@@ -52,9 +56,21 @@ class LSTMModel(FairseqModel):
     @classmethod
     def build_model(cls, args, src_dict, dst_dict):
         """Build a new model instance."""
+
+        encoder_embed_dict = None
+        if args.encoder_embed_path:
+            encoder_embed_dict = utils.parse_embedding(args.encoder_embed_path)
+            utils.print_embed_overlap(encoder_embed_dict, src_dict)
+
+        decoder_embed_dict = None
+        if args.decoder_embed_path:
+            decoder_embed_dict = utils.parse_embedding(args.decoder_embed_path)
+            utils.print_embed_overlap(decoder_embed_dict, dst_dict)
+
         encoder = LSTMEncoder(
             src_dict,
             embed_dim=args.encoder_embed_dim,
+            embed_dict=encoder_embed_dict,
             num_layers=args.encoder_layers,
             dropout_in=args.encoder_dropout_in,
             dropout_out=args.encoder_dropout_out,
@@ -63,6 +79,7 @@ class LSTMModel(FairseqModel):
             dst_dict,
             encoder_embed_dim=args.encoder_embed_dim,
             embed_dim=args.decoder_embed_dim,
+            embed_dict=decoder_embed_dict,
             out_embed_dim=args.decoder_out_embed_dim,
             num_layers=args.decoder_layers,
             attention=bool(eval(args.decoder_attention)),
@@ -74,8 +91,8 @@ class LSTMModel(FairseqModel):
 
 class LSTMEncoder(FairseqEncoder):
     """LSTM encoder."""
-    def __init__(self, dictionary, embed_dim=512, num_layers=1, dropout_in=0.1,
-                 dropout_out=0.1):
+    def __init__(self, dictionary, embed_dim=512, embed_dict=None,
+        num_layers=1, dropout_in=0.1, dropout_out=0.1):
         super().__init__(dictionary)
         self.num_layers = num_layers
         self.dropout_in = dropout_in
@@ -84,6 +101,9 @@ class LSTMEncoder(FairseqEncoder):
         num_embeddings = len(dictionary)
         self.padding_idx = dictionary.pad()
         self.embed_tokens = Embedding(num_embeddings, embed_dim, self.padding_idx)
+        if embed_dict:
+            self.embed_tokens = utils.load_embedding(
+                embed_dict, self.dictionary, self.embed_tokens)
 
         self.lstm = LSTM(
             input_size=embed_dim,
@@ -163,7 +183,8 @@ class AttentionLayer(nn.Module):
 
 class LSTMDecoder(FairseqIncrementalDecoder):
     """LSTM decoder."""
-    def __init__(self, dictionary, encoder_embed_dim=512, embed_dim=512,
+    def __init__(self, dictionary, encoder_embed_dim=512, 
+                 embed_dim=512, embed_dict=None,
                  out_embed_dim=512, num_layers=1, dropout_in=0.1,
                  dropout_out=0.1, attention=True):
         super().__init__(dictionary)
@@ -173,6 +194,10 @@ class LSTMDecoder(FairseqIncrementalDecoder):
         num_embeddings = len(dictionary)
         padding_idx = dictionary.pad()
         self.embed_tokens = Embedding(num_embeddings, embed_dim, padding_idx)
+        if embed_dict:
+            self.embed_tokens = utils.load_embedding(
+                embed_dict, self.dictionary, self.embed_tokens)
+
 
         self.layers = nn.ModuleList([
             LSTMCell(encoder_embed_dim + embed_dim if layer == 0 else embed_dim, embed_dim)

fairseq/utils.py

@@ -248,6 +248,38 @@ def load_align_dict(replace_unk):
     return align_dict
 
 
+def print_embed_overlap(embed_dict, vocab_dict):
+     embed_keys = set(embed_dict.keys())
+     vocab_keys = set(vocab_dict.symbols)
+     overlap = len(embed_keys & vocab_keys)
+     print("| Found {}/{} types in embedding file.".format(overlap, len(vocab_dict)))
+
+def parse_embedding(embed_path):
+    """Parse embedding text file into a dictionary of word and embedding tensors.
+
+    The first line can have vocabulary size and dimension. The following lines
+    should contain word and embedding separated by spaces.
+
+    Example:
+        2 5
+        the -0.0230 -0.0264  0.0287  0.0171  0.1403
+        at -0.0395 -0.1286  0.0275  0.0254 -0.0932
+    """
+    embed_dict = dict()
+    with open(embed_path) as f_embed:
+        _ = next(f_embed) #skip header
+        for line in f_embed:
+            pieces = line.strip().split()
+            embed_dict[pieces[0]] = torch.Tensor([float(weight) for weight in pieces[1:]])
+    return embed_dict
+
+def load_embedding(embed_dict, vocab, embedding):
+    for idx in range(len(vocab)):
+        token = vocab[idx]
+        if token in embed_dict:
+            embedding.weight.data[idx] = embed_dict[token]
+    return embedding
+
 def replace_unk(hypo_str, src_str, alignment, align_dict, unk):
     from fairseq import tokenizer
     # Tokens are strings here