using Knet, Random, Distributions

function generate(model, vocab, samples=1; prefix="")
    for n in 1:samples
        reset!(model)
        input = vocab.eos
        for token in split(prefix, "")
            model([input])
            input = vocab.w2i[token]
            print(token)
        end
        while true
            #output = model([input])
            #input = sample(output)
            probs = softmax(vec(model([input])))
            input = rand(Categorical(Array(probs)))
            if input == vocab.eos; break; end
            print(vocab.i2w[input])
        end
        println()
    end
end

function generate_old(model,chars,n)
    x = 1
    reset!(model)
    for i=1:n
        y = model([x])
        x = sample(y)
        print(chars[x])
    end
    println()
end

reset!(m)=(for r in m.layers; r isa RNN && (r.c=r.h=0); end);

function sample(y)
    p = Array(exp.(y)); r = rand()*sum(p)
    for j=1:length(p); (r -= p[j]) < 0 && return j; end
end
