def __init__(self, args, infer=False):

self.args = args

training = not infer

if infer:

args.batch_size = 1

args.seq_length = 1

if args.model == 'rnn':

cell_fn = rnn_cell.BasicRNNCell

elif args.model == 'gru':

cell_fn = rnn_cell.GRUCell

elif args.model == 'lstm':

cell_fn = rnn_cell.BasicLSTMCell

else:

raise Exception("model type not supported: {}".format(args.model))

cell = cell_fn(args.rnn_size)

if training and args.dropout > 0:

cell = rnn_cell.DropoutWrapper(cell, output_keep_prob=1.0-args.dropout)

self.cell = cell = rnn_cell.MultiRNNCell([cell] * args.num_layers)

self.input_data = tf.placeholder(tf.int32, [args.batch_size, args.seq_length])

self.targets = tf.placeholder(tf.int32, [args.batch_size, args.seq_length])

self.initial_state = cell.zero_state(args.batch_size, tf.float32)

with tf.variable_scope('rnnlm'):

softmax_w = tf.get_variable("softmax_w", [args.rnn_size, args.vocab_size])

softmax_b = tf.get_variable("softmax_b", [args.vocab_size])

with tf.device("/cpu:0"):

self.embedding = tf.get_variable("embedding", [args.vocab_size, args.rnn_size])

inputs = tf.nn.embedding_lookup(self.embedding, self.input_data)

if training and args.dropout > 0:

inputs = tf.nn.dropout(inputs, args.dropout)

inputs = tf.split(1, args.seq_length, inputs)

inputs = [tf.squeeze(input_, [1]) for input_ in inputs]

def loop(prev, _):

prev = tf.matmul(prev, softmax_w) + softmax_b

prev_symbol = tf.stop_gradient(tf.argmax(prev, 1))

return tf.nn.embedding_lookup(self.embedding, prev_symbol)

outputs, last_state = seq2seq.rnn_decoder(inputs, self.initial_state, cell, loop_function=loop if infer else None, scope='rnnlm')

output = tf.reshape(tf.concat(1, outputs), [-1, args.rnn_size])

self.logits = tf.matmul(output, softmax_w) + softmax_b

self.probs = tf.nn.softmax(self.logits)

loss = seq2seq.sequence_loss_by_example([self.logits],

[tf.reshape(self.targets, [-1])],

[tf.ones([args.batch_size * args.seq_length])],

args.vocab_size)

self.cost = tf.reduce_sum(loss) / args.batch_size / args.seq_length

self.final_state = last_state

if not infer:

self.lr = tf.Variable(0.0, trainable=False)

tvars = tf.trainable_variables()

grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars),

args.grad_clip)

optimizer = tf.train.AdamOptimizer(self.lr)

self.train_op = optimizer.apply_gradients(zip(grads, tvars))

def sample(self, sess, chars, vocab, num=200, prime='The ', sampling_type=1, temperature=1., word_level=False):

if word_level:

prime = prime.split()

state = self.cell.zero_state(1, tf.float32).eval()

for char in prime[:-1]:

x = np.zeros((1, 1))

x[0, 0] = vocab[char]

feed = {self.input_data: x, self.initial_state:state}

[state] = sess.run([self.final_state], feed)

def weighted_pick(weights):

t = np.cumsum(weights)

s = np.sum(weights)

return(int(np.searchsorted(t, np.random.rand(1)*s)))

if word_level:

ret = list(prime)

else:

ret = prime

char = prime[-1]

for n in range(num):

x = np.zeros((1, 1))

x[0, 0] = vocab[char]

feed = {self.input_data: x, self.initial_state:state}

self.probs = tf.nn.softmax(tf.div(self.logits, temperature))

[probs, state] = sess.run([self.probs, self.final_state], feed)

p = probs[0]

if sampling_type == 0:

sample = np.argmax(p)

elif sampling_type == 2:

if char == ' ':

sample = weighted_pick(p)

else:

sample = np.argmax(p)

else: # sampling_type == 1 default:

sample = weighted_pick(p)

pred = chars[sample]

if word_level:

ret.append(pred)

else:

ret += pred

char = pred

if word_level:

ret = ' '.join(ret)