Generation from pretrained LM model

Author: xiaoda99

generate.py

@@ -0,0 +1,117 @@
+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)
+
+    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)
+            next_idx = torch.multinomial(lm_probs[:, -1, :], 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:')