Skip to content

Commit e4ed2cd

Browse files
committed
sample: Add output file option
Some dependencies erroneously print diagnostic messages to stdout, mixing them with sampling result. -file option allows workaround this
1 parent 690036d commit e4ed2cd

File tree

2 files changed

+10
-1
lines changed

2 files changed

+10
-1
lines changed

doc/flags.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,3 +57,4 @@ The sampling script `sample.lua` accepts the following command-line flags:
5757
- `-gpu`: The ID of the GPU to use (zero-indexed). Default is 0. Set this to -1 to run in CPU-only mode.
5858
- `-gpu_backend`: The GPU backend to use; either `cuda` or `opencl`. Default is `cuda`.
5959
- `-verbose`: By default just the sampled text is printed to the console. Set this to 1 to also print some diagnostic information.
60+
- `-output`: By default sampled text gets output to standard output, together with diagnostic messages. This option specifies the file to output the result.

sample.lua

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ cmd:option('-temperature', 1)
1313
cmd:option('-gpu', 0)
1414
cmd:option('-gpu_backend', 'cuda')
1515
cmd:option('-verbose', 0)
16+
cmd:option('-output', '-')
1617
local opt = cmd:parse(arg)
1718

1819

@@ -39,4 +40,11 @@ if opt.verbose == 1 then print(msg) end
3940
model:evaluate()
4041

4142
local sample = model:sample(opt)
42-
print(sample)
43+
if opt.output == "-" then
44+
print(sample)
45+
else
46+
require 'io'
47+
local outfile = io.open(opt.output, "w")
48+
outfile:write(sample)
49+
outfile:close()
50+
end

0 commit comments

Comments
 (0)