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

self.args = args

if infer:

args.batch_size = 1

args.seq_length = 1

if args.model == 'rnn':

cell_fn = rnn.BasicRNNCell

elif args.model == 'gru':

cell_fn = rnn.GRUCell

elif args.model == 'lstm':

cell_fn = rnn.BasicLSTMCell

elif args.model == 'nas':

cell_fn = rnn.NASCell

else:

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

with tf.device(args.device):

cells = []

for _ in range(args.num_layers):

cell = cell_fn(args.rnn_size)

if not infer and (args.output_keep_prob < 1.0 or args.input_keep_prob < 1.0):

cell = rnn.DropoutWrapper(

cell, input_keep_prob=args.input_keep_prob, output_keep_prob=args.output_keep_prob

)

cells.append(cell)

self.cell = cell = rnn.MultiRNNCell(cells, state_is_tuple=True)

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

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

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])

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

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

if not infer and args.output_keep_prob:

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

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

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(embedding, prev_symbol)

outputs, last_state = legacy_seq2seq.rnn_decoder(inputs, self.initial_state, cell,

loop_function=loop if infer else None, scope='rnnlm')

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

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

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

loss = legacy_seq2seq.sequence_loss_by_example(

[self.logits],

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

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

args.vocab_size

)

self.word_len = tf.placeholder(tf.int32, shape=[args.batch_size], name='word_lengths')

mask = tf.sequence_mask(self.word_len, args.seq_length, dtype=tf.float32)

mask = tf.reshape(mask, [-1])

loss = tf.multiply(mask, loss)

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

self.final_state = last_state

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)

with tf.name_scope('optimizer'):

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

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

# instrument tensorboard

tf.summary.histogram('logits', self.logits)

tf.summary.histogram('loss', loss)

tf.summary.scalar('train_loss', self.cost)

def smash(self, sess, vocab, word):

p_word = np.full((len(word), len(word)), 0.0)

for i in range(len(word)):

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

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

x[0, 0] = vocab[' ']

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

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

p = probs[0]

oldplog = 0

for j in range(i, len(word)):

p_word[i][j] = oldplog + np.log(p[vocab[word[j]]])

oldplog = p_word[i][j]

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

x[0, 0] = vocab[word[j]]

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

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

p = probs[0]

p_word[i][j] += np.log(p[vocab[' ']])

w = np.full(len(word), -1)

f = p_word[0]

for j in range(5, len(word)):

for i in range(2, j - 2):

if f[i] + p_word[i + 1][j] > f[j]:

f[j] = f[i] + p_word[i + 1][j]

w[j] = int(i)

splitted = []

j = len(word) - 1

i = int(w[j])

while j >= 0:

if i != 0:

splitted.append(word[i+1:j+1])

j = i

i = int(w[j-1])

splitted.reverse()