mmt

#!/usr/bin/env python
import argparse
import os
import shutil
import sys

from cli import IllegalStateException, IllegalArgumentException, CorpusNotFoundInFolderException
from cli.evaluation import Evaluator
from cli.libs import osutils, nvidia_smi
from cli.libs.progressbar import Progressbar
from cli.mmt import BilingualCorpus, TMXCorpus, FileParallelCorpus
from cli.mmt.cluster import ClusterNode
from cli.mmt.engine import Engine, EngineBuilder
from cli.translators import MMTTranslator, XLIFFTranslator, BatchTranslator, InteractiveTranslator

__author__ = 'Davide Caroselli and Andrea Rossi'
__description = '''\
  MMT is a context-aware, incremental and general purpose Machine Translation technology.

  MMT goal is to make MT easy to adopt and scale.

  With MMT you don\'t need anymore to train multiple custom engines,
  you can push all your data to a single engine that will automatically
  and in real-time adapt to the context you provide.
  MMT aims to deliver the quality of a custom engine and the low sparsity
  of your all data combined.

  You can find more information on: http://www.modernmt.eu/
'''


def __check_java():
    try:
        _, stderr = osutils.shell_exec(['java', '-version'])

        ok = False
        for line in stderr.split('\n'):
            tokens = line.split()
            if 'version' in tokens:
                if '"1.8' in tokens[tokens.index('version') + 1]:
                    ok = True
                    break
        if not ok:
            print 'ERROR: Wrong version of Java, required Java 8'
            exit(1)
    except OSError:
        print 'ERROR: Missing Java executable, please check INSTALL.md'
        exit(1)


class CLIArgsException(Exception):
    def __init__(self, parser, error):
        self.parser = parser
        self.message = error


def main_create(argv):
    parser = argparse.ArgumentParser(description='Create a new MMT engine from the input corpora')
    parser.prog = 'mmt create'
    parser.add_argument('source_lang', metavar='SOURCE_LANGUAGE', help='the source language (ISO 639-1)')
    parser.add_argument('target_lang', metavar='TARGET_LANGUAGE', help='the target language (ISO 639-1)')
    parser.add_argument('corpora_paths', metavar='CORPORA', nargs='+',
                        help='the paths to the training corpora (you can specify more than one path)')
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('-d', '--debug', action='store_true', dest='debug',
                        help='if debug is set, it enables verbose logging and prevents temporary files to be removed '
                             'after execution')
    parser.add_argument('-s', '--steps', metavar='STEPS', dest='train_steps', choices=EngineBuilder.all_visible_steps(),
                        nargs='+', help='run only specified training steps (possible values are: %(choices)s)')
    parser.add_argument('-r', '--resume', action='store_true', dest='resume', default=False,
                        help='if resume is set, it restores a previously interrupted training')
    parser.add_argument('-y', '--yes', action='store_true', dest='force_delete', default=False,
                        help='if set, skip engine overwrite confirmation check')
    parser.add_argument('--no-split', action='store_false', dest='split_corpora', default=True,
                        help='if no-split is set, MMT will not extract dev and test sets out of the provided '
                             'training corpora')
    parser.add_argument('--validation-corpora', dest='validation_corpora', default=None,
                        help='custom validation corpora path (dev set used by default).')
    parser.add_argument('--batch-size', dest='batch_size', type=int, default=3072,
                        help='batch size used during training (default is 3072)')
    parser.add_argument('--train-steps', dest='n_train_steps', type=int, default=None,
                        help='maximum number of training steps (default is unlimited)')
    parser.add_argument('--eval-steps', dest='eval_freq', type=int, default=1000,
                        help='evaluation frequency during the training (default is 1000)')
    parser.add_argument('--hparams', dest='hparams', default='transformer_base',
                        help='HParams set to use for the training (default is "transformer_base")')
    parser.add_argument('--bpe', dest='bpe_symbols', default=2**15, type=int,
                        help='number of symbols for BPE (default is 32768)')
    parser.add_argument('--gpus', dest='gpus', nargs='+', type=int, default=None,
                        help='the list of GPUs available to MMT (default is all available GPUs, '
                             'use "-1" to train on CPU).')

    if len(argv) == 0:
        parser.print_help()
        return

    # Parse args
    args = parser.parse_args(argv)

    if not args.split_corpora and args.validation_corpora is None:
        raise CLIArgsException(parser, 'you must specify a validation set (remove --no-split option or use '
                                       '--validation-corpora option)')

    if args.gpus is None:
        args.gpus = nvidia_smi.list_gpus()
    elif -1 in args.gpus:
        args.gpus = []

    # Stop the node with the given engine name if it is already running
    node = ClusterNode.connect(args.engine, silent=True)

    if node is not None:
        proceed = True

        if not args.force_delete and not args.resume:
            while True:
                resp = raw_input('An engine named "%s" already exists, '
                                 'are you sure you want to overwrite it? [y/N] ' % args.engine)
                resp = resp.lower()
                if len(resp) == 0 or resp == 'n':
                    proceed = False
                    break
                elif resp == 'y':
                    break

        if not proceed:
            print 'Aborted'
            exit(0)
        else:
            node.stop()

    # Initialize a builder to create and train a new engine
    builder = EngineBuilder(args.engine, args.source_lang, args.target_lang, args.corpora_paths, args.gpus,
                            debug=args.debug, steps=args.train_steps, split_train=args.split_corpora,
                            validation_path=args.validation_corpora, batch_size=args.batch_size,
                            n_train_steps=args.n_train_steps, n_eval_steps=args.eval_freq, hparams=args.hparams,
                            bpe_symbols=args.bpe_symbols)

    if args.resume:
        builder.resume()
    else:
        builder.build()


def main_start(argv):
    parser = argparse.ArgumentParser(description='Start a MMT cluster node')
    parser.prog = 'mmt start'
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('-v', '--verbosity', dest='verbosity', help='log verbosity (0 = only severe errors, '
                                                                    '3 = finest logging)', default=None)
    parser.add_argument('-p', '--api-port', dest='api_port', metavar='API_PORT',
                        help='the public REST Api port. (default is 8045)', default=None, type=int)
    parser.add_argument('--cluster-port', dest='cluster_port', metavar='CLUSTER_PORT',
                        help='the network port used internally by the cluster for communication between '
                             'Cluster nodes. (default is 5016)', default=None, type=int)
    parser.add_argument('--datastream-port', dest='datastream_port', metavar='DATASTREAM_PORT',
                        help='the network port used by Datastream, currently implemented with Kafka '
                             '(default is 9092', default=None, type=int)
    parser.add_argument('--db-port', dest='db_port', metavar='DB_PORT',
                        help='the network port used by the DB, currently implemented with Cassandra '
                             '(default is 9042', default=None, type=int)
    parser.add_argument('--join-leader', dest='leader', metavar='NODE_IP', default=None,
                        help='use this option to join this node to an existent cluster. '
                             'NODE is the IP of the remote host to connect to.')
    parser.add_argument('-d', '--remote-debug', action='store_true', dest='remote_debug',
                        help='setting this option allows Java to connect for remote debug '
                             '(intended only for development purpose)')
    parser.add_argument('--log-file', dest='log_file', default=None, help='custom location for node log file')

    # Parse args
    args = parser.parse_args(argv)

    # create a clusterNode for that engine
    node = ClusterNode.connect(args.engine)
    ensure_not_running(node)

    success = False

    try:
        # start the ClusterNode
        print 'Starting MMT engine \'{engine}\'...'.format(engine=args.engine),
        node.start(api_port=args.api_port,
                   cluster_port=args.cluster_port,
                   datastream_port=args.datastream_port,
                   db_port=args.db_port,
                   leader=args.leader,
                   verbosity=args.verbosity,
                   remote_debug=args.remote_debug,
                   log_file=args.log_file)
        node.wait('JOINED')
        print 'OK'

        print 'Loading models...',
        node.wait('RUNNING')
        print 'OK'

        # the node has started
        print
        print "The MMT engine '" + args.engine + "' is ready."
        print

        if node.api is not None:
            print 'You can try the API with:\n' \
                  '\tcurl "%s/translate?q=world&source=en&target=it&context=computer"' \
                  ' | python -mjson.tool\n' % node.api.base_path
        success = True
    except Exception:
        print 'FAIL'
        raise
    finally:
        if not success:
            node.stop()


def main_stop(argv):
    parser = argparse.ArgumentParser(description='Stop the local instance of MMT engine')
    parser.prog = 'mmt stop'
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('-f', '--forced', action='store_true', dest='forced', default=False,
                        help='forced stop. By default ModernMT will run a graceful shutdown trying to minimize the '
                             'impact on the operations (i.e. fulfilling all pending translation requests before '
                             'stopping the system). If this flag is specified the system will be forcefully halted.')

    # Parse args
    args = parser.parse_args(argv)
    # connect to the already active cluster node
    node = ClusterNode.connect(args.engine)

    print
    print 'Stopping MMT engine \'{engine}\'...'.format(engine=node.engine.name),
    node.stop(force=args.forced)
    print 'OK \n'


def main_status(argv):
    parser = argparse.ArgumentParser(description='Show the MMT engines status')
    parser.prog = 'mmt status'
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default=None)

    # Parse args
    args = parser.parse_args(argv)
    if args.engine is None:
        engines = Engine.list()
    else:
        engines = [args.engine]
    if len(engines) == 0:
        print 'No engine could be found.'
        print 'You can create a new engine with the ./mmt create command.'

    # Get engine names and for each engine connect to its Cluster node and print its state
    for engine_name in engines:
        node = ClusterNode.connect(engine_name)
        node_running = node.running
        node_state = node.state

        rest_api_s = ('running - %s/translate' % node_state.api_port) \
            if node_running else 'stopped'
        cluster_s = ('running - port %d' % node_state.cluster_port) \
            if node_running else 'stopped'
        datastream_s = ('running - %s:%d' % (node_state.datastream_host, node_state.datastream_port)) \
            if node_running else 'stopped'
        database_s = ('running - %s:%d' % (node_state.database_host, node_state.database_port)) \
            if node_running else 'stopped'

        print '[Engine: "%s"]' % engine_name
        print '    REST API:   %s' % rest_api_s
        print '    Cluster:    %s' % cluster_s
        print '    Datastream: %s' % datastream_s
        print '    Database:   %s' % database_s


def main_delete(argv):
    parser = argparse.ArgumentParser(description='Deletes an MMT engine')
    parser.prog = 'mmt delete'
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('--yes', action='store_false', dest='ask_confirmation', default=True,
                        help='if "--yes" is set, this command won\'t ask for confirmation')

    # Parse args
    args = parser.parse_args(argv)

    # connect to the already active cluster node
    node = ClusterNode.connect(args.engine)

    delete = True

    if args.ask_confirmation:
        valid = {'yes': True, 'y': True, 'ye': True, 'no': False, 'n': False}

        while True:
            print 'Are you sure you want to delete engine "%s"? [y/N] ' % args.engine,
            choice = raw_input().lower()

            if choice == '':
                delete = False
                break
            elif choice in valid:
                delete = valid[choice]
                break
            else:
                print 'Please respond with "yes" or "no" (or "y" or "n").'

    if delete:
        print '\nDeleting engine "{engine}"...'.format(engine=args.engine),
        node.stop()
        shutil.rmtree(node.engine.path, ignore_errors=True)
        print 'OK\n'
    else:
        print 'Aborted'


def main_evaluate(argv):
    parser = argparse.ArgumentParser(description='Evaluate MMT engine')
    parser.prog = 'mmt evaluate'
    parser.add_argument('-s', '--source', dest='source_lang', metavar='SOURCE_LANGUAGE', default=None,
                        help='the source language (ISO 639-1). Can be omitted if engine is monolingual.')
    parser.add_argument('-t', '--target', dest='target_lang', metavar='TARGET_LANGUAGE', default=None,
                        help='the target language (ISO 639-1). Can be omitted if engine is monolingual.')
    parser.add_argument('--path', dest='corpora_path', metavar='CORPORA', default=None,
                        help='the path to the test corpora (default is the automatically splitted sample)')
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('--gt-key', dest='gt_key', metavar='GT_API_KEY', default=None,
                        help='A custom Google Translate API Key to use during evaluation')
    parser.add_argument('--human-eval', dest='heval_output', metavar='OUTPUT', default=None,
                        help='the output folder for the tab-spaced files needed to setup a Human Evaluation benchmark')
    parser.add_argument('-d', '--debug', action='store_true', dest='debug', help='if debug is set, it enables verbose '
                                                                                 'logging and prevents temporary files '
                                                                                 'to be removed after execution')

    # Parse args
    args = parser.parse_args(argv)

    # connect to the already active cluster node
    node = ClusterNode.connect(args.engine)
    ensure_running(node)
    ensure_api(node)

    # Infer default arguments
    corpora_path = args.corpora_path or node.engine.test_data_path

    source_lang, target_lang = (args.source_lang, args.target_lang)
    if source_lang is None or target_lang is None:
        if len(node.engine.languages) > 1:
            raise IllegalArgumentException('Missing language. '
                                           'Options "-s" and "-t" are mandatory for multilingual engines.')
        source_lang, target_lang = node.engine.languages[0]

    # perform evaluation
    evaluator = Evaluator(node, source_lang=source_lang, target_lang=target_lang, google_key=args.gt_key)
    corpora = BilingualCorpus.list(source_lang, target_lang, corpora_path)
    evaluator.evaluate(corpora=corpora, heval_output=args.heval_output, debug=args.debug)

    if args.heval_output is not None:
        print 'Files for Human Evaluation are available here:', os.path.abspath(args.heval_output)
        print


def main_translate(argv):
    parser = argparse.ArgumentParser(description='Translate text with ModernMT')

    parser.add_argument('text', metavar='TEXT', help='text to be translated (optional)', default=None, nargs='?')
    parser.add_argument('-s', '--source', dest='source_lang', metavar='SOURCE_LANGUAGE', default=None,
                        help='the source language (ISO 639-1). Can be omitted if engine is monolingual.')
    parser.add_argument('-t', '--target', dest='target_lang', metavar='TARGET_LANGUAGE', default=None,
                        help='the target language (ISO 639-1). Can be omitted if engine is monolingual.')

    # Context arguments
    parser.add_argument('--context', metavar='CONTEXT', dest='context',
                        help='A string to be used as translation context')
    parser.add_argument('--context-file', metavar='CONTEXT_FILE', dest='context_file',
                        help='A local file to be used as translation context')
    parser.add_argument('--context-vector', metavar='CONTEXT_VECTOR', dest='context_vector',
                        help='The context vector with format: <document 1>:<score 1>[,<document N>:<score N>]')

    # Mixed arguments
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('--batch', action='store_true', dest='batch', default=False,
                        help='if set, the script will read the whole stdin before send translations to MMT.'
                             'This can be used to execute translation in parallel for a faster translation. ')
    parser.add_argument('--xliff', dest='is_xliff', action='store_true', default=False,
                        help='if set, the input is a XLIFF file.')

    # Parse args
    args = parser.parse_args(argv)

    # connect to the already active cluster node
    node = ClusterNode.connect(args.engine)
    ensure_running(node)
    ensure_api(node)

    # Infer default arguments
    source_lang, target_lang = (args.source_lang, args.target_lang)
    if source_lang is None or target_lang is None:
        if len(node.engine.languages) > 1:
            raise IllegalArgumentException('Missing language. '
                                           'Options "-s" and "-t" are mandatory for multilingual engines.')
        source_lang, target_lang = node.engine.languages[0]

    # translate
    mmt_translator = MMTTranslator(node, source_lang, target_lang, context_string=args.context,
                                   context_file=args.context_file, context_vector=args.context_vector)

    if args.text is not None:
        print mmt_translator.translate_text(args.text.strip()).encode('utf-8')
    else:
        if args.is_xliff:
            translator = XLIFFTranslator(mmt_translator)
        elif args.batch:
            translator = BatchTranslator(mmt_translator)
        else:
            translator = InteractiveTranslator(mmt_translator)

        try:
            translator.run(sys.stdin, sys.stdout)
        except KeyboardInterrupt:
            pass  # exit


def main_add(argv):
    parser = argparse.ArgumentParser(description='Add contribution to an existent memory')

    parser.add_argument('memory', help='The id or name of the memory you want to add the contribution to')
    parser.add_argument('source', metavar='SOURCE_SENTENCE', help='The source sentence of the contribution')
    parser.add_argument('target', metavar='TARGET_SENTENCE', help='The target sentence of the contribution')
    parser.add_argument('-s', '--source', dest='source_lang', metavar='SOURCE_LANGUAGE', default=None,
                        help='the source language (ISO 639-1). Can be omitted if engine is monolingual.')
    parser.add_argument('-t', '--target', dest='target_lang', metavar='TARGET_LANGUAGE', default=None,
                        help='the target language (ISO 639-1). Can be omitted if engine is monolingual.')
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')

    args = parser.parse_args(argv)

    node = ClusterNode.connect(args.engine)
    ensure_running(node)
    ensure_api(node)

    # Infer default arguments
    source_lang, target_lang = (args.source_lang, args.target_lang)
    if source_lang is None or target_lang is None:
        if len(node.engine.languages) > 1:
            raise IllegalArgumentException('Missing language. '
                                           'Options "-s" and "-t" are mandatory for multilingual engines.')
        source_lang, target_lang = node.engine.languages[0]

    result = node.append_to_memory(args.memory, source_lang, target_lang, args.source, args.target)
    if result is None:
        memory = node.api.create_memory(args.memory)
        print 'Created new memory with name "%s"' % args.memory

        node.append_to_memory(memory['id'], source_lang, target_lang, args.source, args.target)

    print 'SUCCESS - contribution added to memory "' + args.memory + '"'


def main_rename(argv):
    parser = argparse.ArgumentParser(description='Rename an existent memory')

    parser.add_argument('memory', help='The id or name of the memory you want to rename')
    parser.add_argument('name', help='The new name')

    # Mixed arguments
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')

    args = parser.parse_args(argv)

    node = ClusterNode.connect(args.engine)
    ensure_running(node)
    ensure_api(node)

    memory = node.rename_memory(args.memory, args.name)

    print 'SUCCESS - changed memory name to "%s"' % memory['name']


def main_import(argv):
    parser = argparse.ArgumentParser(description='Import a new memory given a TMX')

    parser.add_argument('-x', '--tmx-file', dest='tmx', metavar='TMX_FILE', help='TMX file to import', default=None)
    parser.add_argument('-p', '--parallel-files', dest='parallel_file', help='Source and target file',
                        default=None, nargs=2)
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')

    args = parser.parse_args(argv)

    if args.tmx is None and args.parallel_file is None:
        raise CLIArgsException(parser, 'missing one of the following options: "-x" or "-p"')

    node = ClusterNode.connect(args.engine)
    ensure_running(node)
    ensure_api(node)

    if args.tmx is not None:
        corpus = TMXCorpus.from_file(args.tmx)
    else:
        corpus = FileParallelCorpus.from_files(args.parallel_file[0], args.parallel_file[1])

    memory = node.new_memory(corpus.name)
    progressbar = Progressbar(label='Importing %s' % corpus.name)
    progressbar.start()

    try:
        node.import_corpus(memory['id'], corpus, callback=lambda job: progressbar.set_progress(job['progress']))
        progressbar.complete()
        print 'IMPORT SUCCESS'
    except BaseException as e:
        node.delete_memory(memory['id'])
        progressbar.abort(repr(e))
        print 'IMPORT FAILED'
        raise


def main():
    actions = {
        'create': main_create,
        'start': main_start,
        'stop': main_stop,
        'status': main_status,
        'delete': main_delete,
        'evaluate': main_evaluate,
        'translate': main_translate,
        'add': main_add,
        'rename': main_rename,
        'import': main_import,
    }

    # Set unbuffered stdout
    unbuffered = os.fdopen(sys.stdout.fileno(), 'w', 0)
    sys.stdout = unbuffered

    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, description=__description,
                                     usage='%(prog)s [-h] ACTION [args]', add_help=False, prog='mmt')
    parser.add_argument('action', metavar='ACTION', choices=actions.keys(), help='{%(choices)s}', nargs='?')
    parser.add_argument('-h', '--help', dest='help', action='store_true', help='show this help message and exit')

    argv = sys.argv[1:]

    if len(argv) == 0:
        parser.print_help()
        exit(1)

    command = argv[0]
    args = argv[1:]

    try:
        if command in actions:
            actions[command](args)
        else:
            parser.print_help()
            exit(1)
    except CLIArgsException as e:
        message = '{prog}: error: {message}\n'.format(prog=e.parser.prog, message=e.message)
        e.parser.print_usage(file=sys.stderr)
        sys.stderr.write(message)
        exit(1)
    except CorpusNotFoundInFolderException as e:
        sys.stderr.write('\nERROR Corpus not found: {message}\n'.format(message=e.message))
        exit(1)
    except IllegalArgumentException as e:
        sys.stderr.write('\nERROR Illegal Argument: {message}\n'.format(message=e.message))
        exit(1)
    except IllegalStateException as e:
        sys.stderr.write('\nERROR Illegal State: {message}\n'.format(message=e.message))
        exit(1)
    except KeyboardInterrupt:
        sys.stderr.write('\nERROR Process Interrupted')
        exit(1)
    except Exception as e:
        sys.stderr.write('\nERROR Unexpected exception: {message}\n'.format(message=e.message))
        raise


# Check that the node is running
# and throw an IllegalStateException if it is not
def ensure_running(node):
    if not node.running:
        raise IllegalStateException('MMT engine \'%s\' is not running.\n'
                                    'Start it with "./mmt start"\n'
                                    'You can check the status of your engines with "./mmt status"'
                                    % node.engine.name)


# Check that the node is not running
# and throw an IllegalStateException if it is
def ensure_not_running(node):
    if node.running:
        raise IllegalStateException('MMT engine \'%s\' is already running.\n'
                                    'You can check the status of your engines with "./mmt status"'
                                    % node.engine.name)


# Check that the REST Server of the engine is running
# and throw an IllegalStateException if it is not
def ensure_api(node):
    if node.api is None:
        raise IllegalStateException('No MMT REST Server running. Enable it and restart the engine to perform tuning')


if __name__ == '__main__':
    __check_java()
    main()