edit seed option, add quiet option

Author: graykode

main.py

@@ -6,6 +6,7 @@
 import os
 import sys
 import torch
+import random
 import argparse
 import numpy as np
 from GPT2.model import (GPT2LMHeadModel)
@@ -17,7 +18,7 @@
 def text_generator(state_dict):
     parser = argparse.ArgumentParser()
     parser.add_argument("--text", type=str, required=True)
-    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--quiet", type=bool, default=False)
     parser.add_argument("--nsamples", type=int, default=1)
     parser.add_argument('--unconditional', action='store_true', help='If true, unconditional generation.')
     parser.add_argument("--batch_size", type=int, default=-1)
@@ -31,9 +32,10 @@ def text_generator(state_dict):
         args.batch_size = 1
     assert args.nsamples % args.batch_size == 0
 
-    np.random.seed(args.seed)
-    torch.random.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
+    seed = random.randint(0, 10000)
+    np.random.seed(seed)
+    torch.random.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     # Load Model
@@ -51,7 +53,7 @@ def text_generator(state_dict):
 
     print(args.text)
     context_tokens = enc.encode(args.text)
-    print(context_tokens)
+
     generated = 0
     for _ in range(args.nsamples // args.batch_size):
         out = sample_sequence(
@@ -65,10 +67,9 @@ def text_generator(state_dict):
         for i in range(args.batch_size):
             generated += 1
             text = enc.decode(out[i])
-            print("=" * 40 + " SAMPLE " + str(generated) + " " + "=" * 40)
+            if args.quiet is False:
+                print("=" * 40 + " SAMPLE " + str(generated) + " " + "=" * 40)
             print(text)
-    print("=" * 80)
-
 
 if __name__ == '__main__':
     if os.path.exists('gpt2-pytorch_model.bin'):