first commit

Author: maru0kun

README.md

@@ -0,0 +1,31 @@
+# Effective Approaches to Attention-based Neural Machine Translation
+Encoder-Decoder model with global attention mechanismのpytorch実装。
+
+## Model Details
+- LSTM-based encoder-decoder model
+- global attention (see Figure 2 in original paper)
+- scheduled sampling
+
+
+## Usages
+学習
+```python
+python train.py \
+    --gpu
+    --train ./sample_data/sample_train.py
+    --valid ./sample_data/sample_valid.py
+    --tf-ratio 0.5
+    --savedir ./checkpoints
+```
+
+翻訳
+```python
+python translate.py \
+    --gpu
+    --model ./checkpoints/checkpoint_best.pt
+    --input ./sample_data/sample_test.txt
+```
+
+## References
+- [Effective Approaches to Attention-based Neural Machine Translation](https://arxiv.org/pdf/1508.04025.pdf)
+- [Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks](https://arxiv.org/pdf/1506.03099.pdf)

model.py

@@ -0,0 +1,130 @@
+# -*- coding: utf-8 -*-
+
+import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class EncRNN(nn.Module):
+    def __init__(self, vsz, embed_dim, hidden_dim, n_layers, use_birnn, dout):
+        super(EncRNN, self).__init__()
+        self.embed = nn.Embedding(vsz, embed_dim)
+        self.rnn = nn.LSTM(embed_dim, hidden_dim, n_layers,
+                           bidirectional=use_birnn)
+        self.dropout = nn.Dropout(dout)
+
+    def forward(self, inputs):
+        embs = self.dropout(self.embed(inputs))
+        enc_outs, hidden = self.rnn(embs)
+        return self.dropout(enc_outs), hidden
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden_dim, method):
+        super(Attention, self).__init__()
+        self.method = method
+        self.hidden_dim = hidden_dim
+
+        if method == 'general':
+            self.w = nn.Linear(hidden_dim, hidden_dim)
+        elif method == 'concat':
+            self.w = nn.Linear(hidden_dim*2, hidden_dim)
+            self.v = torch.nn.Parameter(torch.FloatTensor(hidden_dim))
+
+    def forward(self, dec_out, enc_outs):
+        if self.method == 'dot':
+            attn_energies = self.dot(dec_out, enc_outs)
+        elif self.method == 'general':
+            attn_energies = self.general(dec_out, enc_outs)
+        elif self.method == 'concat':
+            attn_energies = self.concat(dec_out, enc_outs)
+        return F.softmax(attn_energies, dim=0)
+
+    def dot(self, dec_out, enc_outs):
+        return torch.sum(dec_out*enc_outs, dim=2)
+
+    def general(self, dec_out, enc_outs):
+        energy = self.w(enc_outs)
+        return torch.sum(dec_out*energy, dim=2)
+
+    def concat(self, dec_out, enc_outs):
+        dec_out = dec_out.expand(enc_outs.shape[0], -1, -1)
+        energy = torch.cat((dec_out, enc_outs), 2)
+        return torch.sum(self.v * self.w(energy).tanh(), dim=2)
+
+
+class DecRNN(nn.Module):
+    def __init__(self, vsz, embed_dim, hidden_dim, n_layers, use_birnn, 
+                 dout, attn, tied):
+        super(DecRNN, self).__init__()
+        hidden_dim = hidden_dim*2 if use_birnn else hidden_dim
+
+        self.embed = nn.Embedding(vsz, embed_dim)
+        self.rnn = nn.LSTM(embed_dim, hidden_dim , n_layers)
+
+        self.w = nn.Linear(hidden_dim*2, hidden_dim)
+        self.attn = Attention(hidden_dim, attn)
+
+        self.out_projection = nn.Linear(hidden_dim, vsz)
+        if tied: 
+            if embed_dim != hidden_dim:
+                raise ValueError(
+                    f"when using the tied flag, embed-dim:{embed_dim} \
+                    must be equal to hidden-dim:{hidden_dim}")
+            self.out_projection.weight = self.embed.weight
+        self.dropout = nn.Dropout(dout)
+
+    def forward(self, inputs, hidden, enc_outs):
+        inputs = inputs.unsqueeze(0)
+        embs = self.dropout(self.embed(inputs))
+        dec_out, hidden = self.rnn(embs, hidden)
+
+        attn_weights = self.attn(dec_out, enc_outs).transpose(1, 0)
+        enc_outs = enc_outs.transpose(1, 0)
+        context = torch.bmm(attn_weights.unsqueeze(1), enc_outs)
+        cats = self.w(torch.cat((dec_out, context.transpose(1, 0)), dim=2))
+        pred = self.out_projection(cats.tanh().squeeze(0))
+        return pred, hidden
+
+
+class Seq2seqAttn(nn.Module):
+    def __init__(self, args, fields, device):
+        super().__init__()
+        self.src_field, self.tgt_field = fields
+        self.src_vsz = len(self.src_field[1].vocab.itos)
+        self.tgt_vsz = len(self.tgt_field[1].vocab.itos)
+        self.encoder = EncRNN(self.src_vsz, args.embed_dim, args.hidden_dim, 
+                              args.n_layers, args.bidirectional, args.dropout)
+        self.decoder = DecRNN(self.tgt_vsz, args.embed_dim, args.hidden_dim, 
+                              args.n_layers, args.bidirectional, args.dropout,
+                              args.attn, args.tied)
+        self.device = device
+        self.n_layers = args.n_layers
+        self.hidden_dim = args.hidden_dim
+        self.use_birnn = args.bidirectional
+
+    def forward(self, srcs, tgts=None, maxlen=100, tf_ratio=0.0):
+        slen, bsz = srcs.size()
+        tlen = tgts.size(0) if isinstance(tgts, torch.Tensor) else maxlen
+        tf_ratio = tf_ratio if isinstance(tgts, torch.Tensor) else 0.0
+       
+        enc_outs, hidden = self.encoder(srcs)
+
+        dec_inputs = torch.ones_like(srcs[0]) * 2 # <eos> is mapped to id=2
+        outs = []
+
+        if self.use_birnn:
+            def trans_hidden(hs):
+                hs = hs.view(self.n_layers, 2, bsz, self.hidden_dim)
+                hs = torch.stack([torch.cat((h[0], h[1]), 1) for h in hs])
+                return hs
+            hidden = tuple(trans_hidden(hs) for hs in hidden)
+
+        for i in range(tlen):
+            preds, hidden = self.decoder(dec_inputs, hidden, enc_outs)
+            outs.append(preds)
+            use_tf = random.random() < tf_ratio
+            dec_inputs = tgts[i] if use_tf else preds.max(1)[1]
+        return torch.stack(outs)

options.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+
+
+def train_opts(parser):
+    group = parser.add_argument_group('Training')
+    group.add_argument('--train', default='./sample_data/sample_train.tsv',
+        help='path to a train data')
+    group.add_argument('--valid', default='./sample_data/sample_valid.tsv',
+        help='path to a validation data')
+    group.add_argument('--batch-size', type=int, default=32, 
+        help='batch size')
+    group.add_argument('--savedir', default='./checkpoints', 
+        help='path to save models')
+    group.add_argument('--max-epoch', type=int, default=0, 
+        help='number of epochs')
+    group.add_argument('--max-update', type=int, default=0,
+        help='number of updates')
+    group.add_argument('--lr', type=float, default=0.25,
+        help='learning rate')
+    group.add_argument('--min-lr', type=float, default=1e-5, 
+        help='minimum learning rate')
+    group.add_argument('--clip', type=float, default=0.1,
+        help='gradient cliping')
+    group.add_argument('--tf-ratio', type=float, default=0.5,
+        help='teaching force ratio')
+    group.add_argument('--gpu', action='store_true',
+        help='whether gpu is used')
+    return group
+
+
+def translate_opts(parser):
+    group = parser.add_argument_group('Translation')
+    group.add_argument('--model', default='./checkpoints/checkpoint_best.pt',
+        help='model file for translation')
+    group.add_argument('--input', default='./sample_data/sample_test.txt',
+        help='input file')
+    group.add_argument('--batch-size', type=int, default=32,
+        help='batch size')
+    group.add_argument('--maxlen', type=int, default=100,
+        help='maximum length of output sentence')
+    group.add_argument('--gpu', action='store_true',
+        help='whether gpu is used')
+    return group
+    
+
+def model_opts(parser):
+    group = parser.add_argument_group('Model\'s hyper-parameters')
+    group.add_argument('--embed-dim', type=int, default=200,
+        help='dimension of word embeddings')
+    group.add_argument('--src_min-freq', type=int, default=0,
+        help='''map words of source side appearing less than 
+                threshold times to unknown''')
+    group.add_argument('--tgt_min-freq', type=int, default=0,
+        help='''map words of target side appearing less than
+              threshold times to unknown''')
+    group.add_argument('--rnn', choices=['lstm'], default='lstm',
+        help='rnn\'s architechture')
+    group.add_argument('--hidden-dim', type=int, default=1024,
+        help='number of hidden units per layer')
+    group.add_argument('--n-layers', type=int, default=2,
+        help='number of LSTM layers')
+    group.add_argument('--bidirectional', action='store_true',
+        help='whether use bidirectional LSTM for encoder')
+    group.add_argument('--attn', choices=['dot', 'general', 'concat'],
+        default='dot', help='attention type')
+    group.add_argument('--dropout', type=float, default=0.2,
+        help='dropout applied to layers (0 means no dropout)')
+    group.add_argument('--tied', action='store_true',
+        help='tie the word embedding and softmax weight')
+    return group
+