make sample.py produce n sentences

carrotflakes · Sep 12, 2017 · 814df0d · 814df0d
1 parent 0db5a37
commit 814df0d
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/sample.py b/sample.py
@@ -39,6 +39,8 @@
                         help='directory of the trained model')
     parser.add_argument('-c', '--only_cpu', default=True, action='store_true',
                         help='if set, only build weights on cpu')
+    parser.add_argument('-i', '--num', default=10, type=int,
+                        help='number of sample to generate')
     args = parser.parse_args()
 
     if not os.path.exists(args.dictionary):
@@ -56,5 +58,6 @@
     model.build()
     model.load(ignore_missing=True)
 
-    g = model.generate(args.sample_len)
-    print('Generated text:', utils.detokenize(g, rev_dict))
+    for _ in range(args.num):
+        g = model.generate(args.sample_len)
+        print('Generated text:', utils.detokenize(g, rev_dict))