Merge pull request #35 from xiaoda99/master

Text generation with pretrained LM model

Author: thomwolf

generate.py

@@ -0,0 +1,122 @@
+import argparse
+import os
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from model_pytorch import LMModel, load_openai_pretrained_model
+from text_utils import TextEncoder
+
+
+def make_batch(X):
+    X = np.array(X)
+    assert X.ndim in [1, 2]
+    if X.ndim == 1:
+        X = np.expand_dims(X, axis=0)
+    pos_enc = np.arange(n_vocab + n_special, n_vocab + n_special + X.shape[-1])
+    pos_enc = np.expand_dims(pos_enc, axis=0)
+    batch = np.stack([X, pos_enc], axis=-1)
+    batch = torch.tensor(batch, dtype=torch.long).to(device)
+    return batch
+
+def append_batch(X, next_idx):
+    next_pos = X[:, -1:, 1] + 1
+    next_x = torch.cat((next_idx, next_pos), -1).unsqueeze(1)
+    return torch.cat((X, next_x), 1)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--desc', type=str, help="Description")
+    parser.add_argument('--dataset', type=str)
+    parser.add_argument('--log_dir', type=str, default='log/')
+    parser.add_argument('--save_dir', type=str, default='save/')
+    parser.add_argument('--data_dir', type=str, default='data/')
+    parser.add_argument('--submission_dir', type=str, default='submission/')
+    parser.add_argument('--submit', action='store_true')
+    parser.add_argument('--analysis', action='store_true')
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--n_iter', type=int, default=3)
+    parser.add_argument('--n_batch', type=int, default=8)
+    parser.add_argument('--max_grad_norm', type=int, default=1)
+    parser.add_argument('--lr', type=float, default=6.25e-5)
+    parser.add_argument('--lr_warmup', type=float, default=0.002)
+    parser.add_argument('--n_ctx', type=int, default=512)
+    parser.add_argument('--n_embd', type=int, default=768)
+    parser.add_argument('--n_head', type=int, default=12)
+    parser.add_argument('--n_layer', type=int, default=12)
+    parser.add_argument('--embd_pdrop', type=float, default=0.1)
+    parser.add_argument('--attn_pdrop', type=float, default=0.1)
+    parser.add_argument('--resid_pdrop', type=float, default=0.1)
+    parser.add_argument('--clf_pdrop', type=float, default=0.1)
+    parser.add_argument('--l2', type=float, default=0.01)
+    parser.add_argument('--vector_l2', action='store_true')
+    parser.add_argument('--opt', type=str, default='adam')
+    parser.add_argument('--afn', type=str, default='gelu')
+    parser.add_argument('--lr_schedule', type=str, default='warmup_linear')
+    parser.add_argument('--encoder_path', type=str, default='model/encoder_bpe_40000.json')
+    parser.add_argument('--bpe_path', type=str, default='model/vocab_40000.bpe')
+    parser.add_argument('--n_transfer', type=int, default=12)
+    parser.add_argument('--lm_coef', type=float, default=0.5)
+    parser.add_argument('--b1', type=float, default=0.9)
+    parser.add_argument('--b2', type=float, default=0.999)
+    parser.add_argument('--e', type=float, default=1e-8)
+    parser.add_argument('--n_valid', type=int, default=374)
+    parser.add_argument('--gen_len', type=int, default=20)
+    parser.add_argument('--topk', type=int, default=10)
+
+    args = parser.parse_args()
+    print(args)
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+
+    # Constants
+    submit = args.submit
+    dataset = args.dataset
+    n_ctx = args.n_ctx
+    save_dir = args.save_dir
+    desc = args.desc
+    data_dir = args.data_dir
+    log_dir = args.log_dir
+    submission_dir = args.submission_dir
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    print("device", device, "n_gpu", n_gpu)
+
+    text_encoder = TextEncoder(args.encoder_path, args.bpe_path)
+    encoder = text_encoder.encoder
+    n_vocab = len(text_encoder.encoder)
+
+    n_special = 0   # XD: useless for language modeling task
+    vocab = n_vocab + n_special + n_ctx
+
+    lm_model = LMModel(args, vocab, n_ctx, return_probs=True)
+    load_openai_pretrained_model(lm_model.transformer, n_ctx=n_ctx, n_special=n_special)
+    lm_model.to(device)
+
+    lm_model.eval()
+
+    text = input('Input some beginning words:')
+    while text != 'q':
+        X = text_encoder.encode([text,])
+        XMB = make_batch(X)
+
+        for _ in range(args.gen_len):
+            lm_probs = lm_model(XMB)
+            if args.topk == 0:
+                next_idx = torch.multinomial(lm_probs[:, -1, :], 1)
+            else:
+                values, indices = lm_probs[:, -1, :].topk(args.topk)
+                next_idx = indices.gather(-1, torch.multinomial(values, 1))
+            next_token = text_encoder.decoder[next_idx.item()].replace('</w>', '')
+            print(next_token, end=' ')
+            XMB = append_batch(XMB, next_idx)
+
+        print()
+        text = input('Input some beginning words:')

model_pytorch.py

@@ -84,7 +84,11 @@ def _attn(self, q, k, v):
         w = torch.matmul(q, k)
         if self.scale:
             w = w / math.sqrt(v.size(-1))
-        w = w * self.b + -1e9 * (1 - self.b)  # TF implem method: mask_attn_weights
+        # w = w * self.b + -1e9 * (1 - self.b)  # TF implem method: mask_attn_weights
+        # XD: self.b may be larger than w, so we need to crop it
+        b = self.b[:, :, :w.size(-2), :w.size(-1)]
+        w = w * b + -1e9 * (1 - b)
+
         w = nn.Softmax(dim=-1)(w)
         w = self.attn_dropout(w)
         return torch.matmul(w, v)
@@ -173,16 +177,18 @@ def forward(self, x):
 class LMHead(nn.Module):
     """ Language Model Head for the transformer """
 
-    def __init__(self, model, cfg):
+    def __init__(self, model, cfg, trunc_and_reshape=True):
         super(LMHead, self).__init__()
         self.n_embd = cfg.n_embd
         embed_shape = model.embed.weight.shape
         self.decoder = nn.Linear(embed_shape[1], embed_shape[0], bias=False)
         self.decoder.weight = model.embed.weight # Tied weights
+        self.trunc_and_reshape = trunc_and_reshape  # XD
 
     def forward(self, h):
         # Truncated Language modeling logits (we remove the last token)
-        h_trunc = h[:, :-1].contiguous().view(-1, self.n_embd)
+        h_trunc = h[:, :-1].contiguous().view(-1, self.n_embd) \
+            if self.trunc_and_reshape else h  # XD
         lm_logits = self.decoder(h_trunc)
         return lm_logits
 
@@ -264,6 +270,29 @@ def forward(self, h, x):
 
         return sim_logits
 
+
+# XD
+class LMModel(nn.Module):
+    """ Transformer with language model head only """
+    def __init__(self, cfg, vocab=40990, n_ctx=512, return_probs=False):
+        super(LMModel, self).__init__()
+        self.transformer = TransformerModel(cfg, vocab=vocab, n_ctx=n_ctx)
+        self.lm_head = LMHead(self.transformer, cfg, trunc_and_reshape=False)
+        self.return_probs = return_probs
+        if self.return_probs:
+            pos_emb_mask = torch.zeros(1, 1, vocab)
+            pos_emb_mask[:, :, -n_ctx:] = -1e12
+            self.register_buffer('pos_emb_mask', pos_emb_mask)
+
+
+    def forward(self, x):
+        h = self.transformer(x)
+        lm_logits = self.lm_head(h)
+        if self.return_probs:
+            lm_logits = F.softmax(lm_logits + self.pos_emb_mask, dim=-1)
+        return lm_logits
+
+
 class DoubleHeadModel(nn.Module):
     """ Transformer with language model and task specific heads """
     def __init__(self, cfg, clf_token, task_head_type, vocab=40990, n_ctx=512):