# Copyright 2016 Google Inc. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""A program to train a tensorflow neural net parser from a a conll file."""

import os

import os.path

import time

import tensorflow as tf

from tensorflow.python.platform import gfile

from tensorflow.python.platform import tf_logging as logging

from google.protobuf import text_format

from syntaxnet import graph_builder

from syntaxnet import structured_graph_builder

from syntaxnet.ops import gen_parser_ops

from syntaxnet import task_spec_pb2

flags = tf.app.flags

FLAGS = flags.FLAGS

flags.DEFINE_string('tf_master', '',

'TensorFlow execution engine to connect to.')

flags.DEFINE_string('output_path', '', 'Top level for output.')

flags.DEFINE_string('task_context', '',

'Path to a task context with resource locations and '

'parameters.')

flags.DEFINE_string('arg_prefix', None, 'Prefix for context parameters.')

flags.DEFINE_string('params', '0', 'Unique identifier of parameter grid point.')

flags.DEFINE_string('training_corpus', 'training-corpus',

'Name of the context input to read training data from.')

flags.DEFINE_string('tuning_corpus', 'tuning-corpus',

'Name of the context input to read tuning data from.')

flags.DEFINE_string('word_embeddings', None,

'Recordio containing pretrained word embeddings, will be '

'loaded as the first embedding matrix.')

flags.DEFINE_bool('compute_lexicon', False, '')

flags.DEFINE_bool('projectivize_training_set', False, '')

flags.DEFINE_string('hidden_layer_sizes', '200,200',

'Comma separated list of hidden layer sizes.')

flags.DEFINE_string('graph_builder', 'greedy',

'Graph builder to use, either "greedy" or "structured".')

flags.DEFINE_integer('batch_size', 32,

'Number of sentences to process in parallel.')

flags.DEFINE_integer('beam_size', 10, 'Number of slots for beam parsing.')

flags.DEFINE_integer('num_epochs', 10, 'Number of epochs to train for.')

flags.DEFINE_integer('max_steps', 50,

'Max number of parser steps during a training step.')

flags.DEFINE_integer('report_every', 100,

'Report cost and training accuracy every this many steps.')

flags.DEFINE_integer('checkpoint_every', 5000,

'Measure tuning UAS and checkpoint every this many steps.')

flags.DEFINE_bool('slim_model', False,

'Whether to remove non-averaged variables, for compactness.')

flags.DEFINE_float('learning_rate', 0.1, 'Initial learning rate parameter.')

flags.DEFINE_integer('decay_steps', 4000,

'Decay learning rate by 0.96 every this many steps.')

flags.DEFINE_float('momentum', 0.9,

'Momentum parameter for momentum optimizer.')

flags.DEFINE_string('seed', '0', 'Initialization seed for TF variables.')

flags.DEFINE_string('pretrained_params', None,

'Path to model from which to load params.')

flags.DEFINE_string('pretrained_params_names', None,

'List of names of tensors to load from pretrained model.')

flags.DEFINE_float('averaging_decay', 0.9999,

'Decay for exponential moving average when computing'

'averaged parameters, set to 1 to do vanilla averaging.')

def StageName():

return os.path.join(FLAGS.arg_prefix, FLAGS.graph_builder)

def OutputPath(path):

return os.path.join(FLAGS.output_path, StageName(), FLAGS.params, path)

def RewriteContext():

context = task_spec_pb2.TaskSpec()

with gfile.FastGFile(FLAGS.task_context, 'rb') as fin:

text_format.Merge(fin.read(), context)

for resource in context.input:

if resource.creator == StageName():

del resource.part[:]

part = resource.part.add()

part.file_pattern = os.path.join(OutputPath(resource.name))

with gfile.FastGFile(OutputPath('context'), 'w') as fout:

fout.write(str(context))

def WriteStatus(num_steps, eval_metric, best_eval_metric):

status = os.path.join(os.getenv('GOOGLE_STATUS_DIR') or '/tmp', 'STATUS')

message = ('Parameters: %s | Steps: %d | Tuning score: %.2f%% | '

'Best tuning score: %.2f%%' % (FLAGS.params, num_steps,

eval_metric, best_eval_metric))

with gfile.FastGFile(status, 'w') as fout:

fout.write(message)

with gfile.FastGFile(OutputPath('status'), 'a') as fout:

fout.write(message + '\n')

def Eval(sess, parser, num_steps, best_eval_metric):

"""Evaluates a network and checkpoints it to disk.

Args:

sess: tensorflow session to use

parser: graph builder containing all ops references

num_steps: number of training steps taken, for logging

best_eval_metric: current best eval metric, to decide whether this model is

the best so far

Returns:

new best eval metric

"""

logging.info('Evaluating training network.')

t = time.time()

num_epochs = None

num_tokens = 0

num_correct = 0

while True:

tf_eval_epochs, tf_eval_metrics = sess.run([

parser.evaluation['epochs'], parser.evaluation['eval_metrics']

])

num_tokens += tf_eval_metrics[0]

num_correct += tf_eval_metrics[1]

if num_epochs is None:

num_epochs = tf_eval_epochs

elif num_epochs < tf_eval_epochs:

break

eval_metric = 0 if num_tokens == 0 else (100.0 * num_correct / num_tokens)

logging.info('Seconds elapsed in evaluation: %.2f, '

'eval metric: %.2f%%', time.time() - t, eval_metric)

WriteStatus(num_steps, eval_metric, max(eval_metric, best_eval_metric))

# Save parameters.

if FLAGS.output_path:

logging.info('Writing out trained parameters.')

parser.saver.save(sess, OutputPath('latest-model'))

if eval_metric > best_eval_metric:

parser.saver.save(sess, OutputPath('model'))

return max(eval_metric, best_eval_metric)

def Train(sess, num_actions, feature_sizes, domain_sizes, embedding_dims):

"""Builds and trains the network.

Args:

sess: tensorflow session to use.

num_actions: number of possible golden actions.

feature_sizes: size of each feature vector.

domain_sizes: number of possible feature ids in each feature vector.

embedding_dims: embedding dimension to use for each feature group.

"""

t = time.time()

hidden_layer_sizes = map(int, FLAGS.hidden_layer_sizes.split(','))

logging.info('Building training network with parameters: feature_sizes: %s '

'domain_sizes: %s', feature_sizes, domain_sizes)

if FLAGS.graph_builder == 'greedy':

parser = graph_builder.GreedyParser(num_actions,

feature_sizes,

domain_sizes,

embedding_dims,

hidden_layer_sizes,

seed=int(FLAGS.seed),

gate_gradients=True,

averaging_decay=FLAGS.averaging_decay,

arg_prefix=FLAGS.arg_prefix)

else:

parser = structured_graph_builder.StructuredGraphBuilder(

num_actions,

feature_sizes,

domain_sizes,

embedding_dims,

hidden_layer_sizes,

seed=int(FLAGS.seed),

gate_gradients=True,

averaging_decay=FLAGS.averaging_decay,

arg_prefix=FLAGS.arg_prefix,

beam_size=FLAGS.beam_size,

max_steps=FLAGS.max_steps)

task_context = OutputPath('context')

if FLAGS.word_embeddings is not None:

parser.AddPretrainedEmbeddings(0, FLAGS.word_embeddings, task_context)

corpus_name = ('projectivized-training-corpus' if

FLAGS.projectivize_training_set else FLAGS.training_corpus)

parser.AddTraining(task_context,

FLAGS.batch_size,

learning_rate=FLAGS.learning_rate,

momentum=FLAGS.momentum,

decay_steps=FLAGS.decay_steps,

corpus_name=corpus_name)

parser.AddEvaluation(task_context,

FLAGS.batch_size,

corpus_name=FLAGS.tuning_corpus)

parser.AddSaver(FLAGS.slim_model)

# Save graph.

if FLAGS.output_path:

with gfile.FastGFile(OutputPath('graph'), 'w') as f:

f.write(sess.graph_def.SerializeToString())

logging.info('Initializing...')

num_epochs = 0

cost_sum = 0.0

num_steps = 0

best_eval_metric = 0.0

sess.run(parser.inits.values())

if FLAGS.pretrained_params is not None:

logging.info('Loading pretrained params from %s', FLAGS.pretrained_params)

feed_dict = {'save/Const:0': FLAGS.pretrained_params}

targets = []

for node in sess.graph_def.node:

if (node.name.startswith('save/Assign') and

node.input[0] in FLAGS.pretrained_params_names.split(',')):

logging.info('Loading %s with op %s', node.input[0], node.name)

targets.append(node.name)

sess.run(targets, feed_dict=feed_dict)

logging.info('Training...')

while num_epochs < FLAGS.num_epochs:

tf_epochs, tf_cost, _ = sess.run([parser.training[

'epochs'], parser.training['cost'], parser.training['train_op']])

num_epochs = tf_epochs

num_steps += 1

cost_sum += tf_cost

if num_steps % FLAGS.report_every == 0:

logging.info('Epochs: %d, num steps: %d, '

'seconds elapsed: %.2f, avg cost: %.2f, ', num_epochs,

num_steps, time.time() - t, cost_sum / FLAGS.report_every)

cost_sum = 0.0

if num_steps % FLAGS.checkpoint_every == 0:

best_eval_metric = Eval(sess, parser, num_steps, best_eval_metric)

def main(unused_argv):

logging.set_verbosity(logging.INFO)

if not gfile.IsDirectory(OutputPath('')):

gfile.MakeDirs(OutputPath(''))

# Rewrite context.

RewriteContext()

# Creates necessary term maps.

if FLAGS.compute_lexicon:

logging.info('Computing lexicon...')

with tf.Session(FLAGS.tf_master) as sess:

gen_parser_ops.lexicon_builder(task_context=OutputPath('context'),

corpus_name=FLAGS.training_corpus).run()

with tf.Session(FLAGS.tf_master) as sess:

feature_sizes, domain_sizes, embedding_dims, num_actions = sess.run(

gen_parser_ops.feature_size(task_context=OutputPath('context'),

arg_prefix=FLAGS.arg_prefix))

# Well formed and projectivize.

if FLAGS.projectivize_training_set:

logging.info('Preprocessing...')

with tf.Session(FLAGS.tf_master) as sess:

source, last = gen_parser_ops.document_source(

task_context=OutputPath('context'),

batch_size=FLAGS.batch_size,

corpus_name=FLAGS.training_corpus)

sink = gen_parser_ops.document_sink(

task_context=OutputPath('context'),

corpus_name='projectivized-training-corpus',

documents=gen_parser_ops.projectivize_filter(

gen_parser_ops.well_formed_filter(source,

task_context=OutputPath(

'context')),

task_context=OutputPath('context')))

while True:

tf_last, _ = sess.run([last, sink])

if tf_last:

break

logging.info('Training...')

with tf.Session(FLAGS.tf_master) as sess:

Train(sess, num_actions, feature_sizes, domain_sizes, embedding_dims)

if __name__ == '__main__':

tf.app.run()