mmt

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

from cli import dependency, IllegalStateException, IllegalArgumentException
from cli.cluster import ClusterNode, DEFAULT_MMT_API_PORT, DEFAULT_MMT_CLUSTER_PORTS, DEFAULT_MMT_DATASTREAM_PORT
from cli.engine import MMTEngine
from cli.evaluation import Evaluator
from cli.libs import netutils
from cli.libs import shell
from cli.mt import BilingualCorpus
from cli.mt.processing import TrainingPreprocessor
from cli.translation import BatchTranslator, InteractiveTranslator, XLIFFTranslator

__author__ = 'Davide Caroselli'
__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 = shell.execute(['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_tune(argv):
    parser = argparse.ArgumentParser(description='Tune MMT engine')
    parser.prog = 'mmt tune'
    parser.add_argument('--path', dest='corpora_path', metavar='CORPORA', default=None,
                        help='the path to the training 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('-p', '--api-port', dest='api_port', metavar='API_PORT',
                        help='the MMT API port. (default is {port})'.format(port=DEFAULT_MMT_API_PORT),
                        default=None, type=int)
    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('--skip-context-analysis', dest='context_enabled', help='if set, context analysis is skipped',
                        default=True, action='store_false')
    parser.add_argument('--random-seeds', dest='random_seeds', help='if set, uses random seed for tuning',
                        default=False, action='store_true')
    parser.add_argument('--max-iterations', dest='max_iterations',
                        help='set maximum iterations during tuning (default is 25)',
                        default=25, type=int)

    args = parser.parse_args(argv)
    args.api_port = args.api_port if args.api_port is not None else DEFAULT_MMT_API_PORT

    engine = MMTEngine.load(args.engine)
    if not engine.exists():
        raise IllegalStateException('engine not found: ' + args.engine)

    corpora = BilingualCorpus.list(args.corpora_path) if args.corpora_path is not None else None

    node = ClusterNode(engine, api_port=args.api_port)
    node.tune(corpora, debug=args.debug, context_enabled=args.context_enabled,
              random_seeds=args.random_seeds, max_iterations=args.max_iterations)


def main_evaluate(argv):
    parser = argparse.ArgumentParser(description='Evaluate MMT engine')
    parser.prog = 'mmt evaluate'
    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('-p', '--api-port', dest='api_port', metavar='API_PORT',
                        help='the MMT API port. (default is {port})'.format(port=DEFAULT_MMT_API_PORT),
                        default=None, type=int)
    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('--gt-nmt', action='store_true', dest='gt_nmt', default=False,
                        help='Use Neural Google Translate API during evaluation, '
                             'you have to specify a valid key with --gt-key')
    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('-i', '--interactive', action='store_true', dest='interactive',
                        help='Test without sessions, as if translating sentence by sentence')
    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')

    args = parser.parse_args(argv)
    args.api_port = args.api_port if args.api_port is not None else DEFAULT_MMT_API_PORT

    engine = MMTEngine.load(args.engine)
    if not engine.exists():
        raise IllegalStateException('engine not found: ' + args.engine)

    node = ClusterNode(engine, api_port=args.api_port)
    evaluator = Evaluator(node, google_key=args.gt_key, google_nmt=args.gt_nmt, use_sessions=not args.interactive)

    corpora = BilingualCorpus.list(args.corpora_path) if args.corpora_path is not None \
        else BilingualCorpus.list(os.path.join(engine.data_path, TrainingPreprocessor.TEST_FOLDER_NAME))

    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_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 {port})'
                        .format(port=DEFAULT_MMT_API_PORT), default=None, type=int)
    parser.add_argument('--cluster-ports', dest='cluster_ports', metavar='CLUSTER_PORT',
                        help='the network ports used internally by the cluster for data exchange between '
                             'Cluster nodes. (default is {port0} and {port1})'
                        .format(port0=DEFAULT_MMT_CLUSTER_PORTS[0], port1=DEFAULT_MMT_CLUSTER_PORTS[1]),
                        default=None, type=int, nargs=2)
    parser.add_argument('--datastream-port', dest='datastream_port', metavar='DATASTREAM_PORT',
                        help='the network port used by Datastream, currently implemented with Kafka '
                             '(default is {port}'.format(port=DEFAULT_MMT_DATASTREAM_PORT),
                        default=None, type=int)
    parser.add_argument('--no-rest', action='store_false', dest='has_rest', default=True,
                        help='if no-rest is set, no REST Api will be available on this node')
    parser.add_argument('--join', dest='sibling', 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.')

    # Parse args
    args = parser.parse_args(argv)

    if not args.has_rest and args.api_port is not None:
        raise CLIArgsException(parser, 'cannot specify api-port while no-rest option is active, '
                                       'please remove option "-p/--api-port"')

    args.api_port = args.api_port if args.api_port is not None else DEFAULT_MMT_API_PORT
    args.cluster_ports = args.cluster_ports if args.cluster_ports is not None else DEFAULT_MMT_CLUSTER_PORTS
    args.datastream_port = args.datastream_port if args.datastream_port is not None else DEFAULT_MMT_DATASTREAM_PORT

    # Check ports
    if args.has_rest and not netutils.is_free(args.api_port):
        raise IllegalStateException(
            'port %d is already in use, please specify another port with option -p/--api-port' % args.api_port)
    if not netutils.is_free(args.cluster_ports[0]):
        raise IllegalStateException(
            'port %d is already in use, please specify another port with --cluster-ports' % args.cluster_ports[0])
    if not netutils.is_free(args.cluster_ports[1]):
        raise IllegalStateException(
            'port %d is already in use, please specify another port with --cluster-ports' % args.cluster_ports[1])

    # Exec command
    engine = MMTEngine.load(args.engine)
    if args.sibling is None and not engine.exists():
        raise IllegalStateException('engine not found: ' + args.engine)

    node = ClusterNode(engine, rest=args.has_rest, api_port=args.api_port, cluster_ports=args.cluster_ports,
                       datastream_port=args.datastream_port, sibling=args.sibling, verbosity=args.verbosity)

    started = False
    success = False

    try:
        print

        print 'Starting MMT engine \'{engine}\'...'.format(engine=engine.name),
        node.start()
        started = True
        node.wait('JOINED')
        print 'OK'

        if args.sibling is not None:
            print 'Synchronizing models...',
            node.wait('SYNCHRONIZED')
            print 'OK'

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

        if args.has_rest is not None:
            print 'You can try the API with:'
            print '\tcurl "http://localhost:{port}/translate?q=world&context=computer"' \
                  ' | python -mjson.tool\n'.format(port=args.api_port)

        success = True
    except Exception:
        print 'FAIL'
        raise
    finally:
        if not success and started and node.is_running():
            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')
    args = parser.parse_args(argv)

    # Exec command
    engine = MMTEngine.load(args.engine)
    if not engine.exists():
        raise IllegalStateException('engine not found: ' + args.engine)

    node = ClusterNode(engine)

    print
    print 'Stopping MMT engine \'{engine}\'...'.format(engine=engine.name),
    if node.is_running():
        node.stop()
    print 'OK'
    print


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')
    args = parser.parse_args(argv)

    # Exec command
    engine = MMTEngine.load(args.engine)
    if not engine.exists():
        raise IllegalStateException('engine not found: ' + 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:
        node = ClusterNode(engine)
        print
        print 'Deleting engine "{engine}"...'.format(engine=engine.name),
        if node.is_running():
            node.stop()
        shutil.rmtree(engine.path, ignore_errors=True)
        print 'OK'
        print
    else:
        print 'Aborted'


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)
    args = parser.parse_args(argv)

    if args.engine is None:
        engines = MMTEngine.list()
    else:
        engines = [MMTEngine.load(args.engine)]
        if not engines[0].exists():
            raise IllegalStateException('engine not found: ' + args.engine)

    for engine in engines:
        node = ClusterNode(engine)
        status = node.get_state()

        print '-'
        print 'Engine:   ', engine.name
        print 'Node:     ', (
            'running - ports %s, %s' % (status['control_port'], status['data_port'])
            if status is not None else 'stopped')
        print 'Rest API: ', (
            'running - http://localhost:%s/translate' % status['api_port']
            if status is not None and 'api_port' in status else 'stopped')


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, '
                             'in every folder you can put mixed monolingual and bilingual corpora')
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default=None)
    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('--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('-s', '--steps', metavar='STEPS', dest='training_steps', choices=MMTEngine.training_steps,
                        nargs='+', help='run only specified training steps {%(choices)s}')

    for component in dependency.injectable_components:
        dependency.argparse_group(parser, component)

    if len(argv) > 0:
        args = parser.parse_args(argv)

        injector = dependency.Injector()
        injector.read_args(args)

        engine = MMTEngine(args.engine, args.source_lang, args.target_lang)
        engine = injector.inject(engine)

        node = ClusterNode(engine)

        if node.is_running():
            node.stop()

        engine.builder.build(args.corpora_paths,
                             debug=args.debug, steps=args.training_steps, split_trainingset=args.split_corpora)
    else:
        parser.print_help()


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='?')

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

    # NBest list arguments
    parser.add_argument('--nbest', metavar='NBEST', dest='nbest', default=None, type=int,
                        help='The number of nbest to print')
    parser.add_argument('--nbest-file', metavar='NBEST_FILE', dest='nbest_file', default=None,
                        help='The destination file for the NBest, default is stdout')

    # Mixed arguments
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('-p', '--api-port', dest='api_port', metavar='API_PORT',
                        help='the public available REST Api port. (default is {port})'
                        .format(port=DEFAULT_MMT_API_PORT), default=None, type=int)
    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.')

    args = parser.parse_args(argv)

    engine = MMTEngine.load(args.engine)
    if not engine.exists():
        raise IllegalStateException('engine not found: ' + args.engine)

    node = ClusterNode(engine, api_port=args.api_port)
    if not node.is_running():
        raise IllegalStateException('engine not running: ' + args.engine)

    if args.is_xliff:
        translator = XLIFFTranslator(node, context_string=args.context, context_file=args.context_file,
                                     context_vector=args.context_vector)
    elif args.batch:
        translator = BatchTranslator(node, context_string=args.context, context_file=args.context_file,
                                     context_vector=args.context_vector, print_nbest=args.nbest,
                                     nbest_file=args.nbest_file)
    else:
        translator = InteractiveTranslator(node, context_string=args.context, context_file=args.context_file,
                                           context_vector=args.context_vector, print_nbest=args.nbest,
                                           nbest_file=args.nbest_file)

    try:
        if args.text is not None:
            translator.execute(args.text.strip())
        else:
            while 1:
                line = sys.stdin.readline()
                if not line:
                    break
                translator.execute(line.strip())

        translator.flush()
    except KeyboardInterrupt:
        # exit
        pass
    finally:
        translator.close()


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

    parser.add_argument('domain', help='The id or name of the domain 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')

    # Mixed arguments
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('-p', '--api-port', dest='api_port', metavar='API_PORT',
                        help='the public available REST Api port. (default is {port})'
                        .format(port=DEFAULT_MMT_API_PORT), default=None, type=int)

    args = parser.parse_args(argv)

    engine = MMTEngine.load(args.engine)
    if not engine.exists():
        raise IllegalStateException('engine not found: ' + args.engine)

    node = ClusterNode(engine, api_port=args.api_port)
    if not node.is_running():
        raise IllegalStateException('engine not running: ' + args.engine)

    node.append_to_domain(args.domain, args.source, args.target)

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


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

    parser.add_argument('-d', '--domain', help='The name of the new domain (default is the filename', default=None)
    parser.add_argument('-x', '--tmx-file', dest='tmx', metavar='TMX_FILE', help='TMX file to import')

    # Mixed arguments
    parser.add_argument('-e', '--engine', dest='engine', help='the engine name, \'default\' will be used if absent',
                        default='default')
    parser.add_argument('-p', '--api-port', dest='api_port', metavar='API_PORT',
                        help='the public available REST Api port. (default is {port})'
                        .format(port=DEFAULT_MMT_API_PORT), default=None, type=int)

    args = parser.parse_args(argv)

    if args.tmx is None:
        raise CLIArgsException(parser, 'missing option "-x"')

    engine = MMTEngine.load(args.engine)
    if not engine.exists():
        raise IllegalStateException('engine not found: ' + args.engine)

    node = ClusterNode(engine, api_port=args.api_port)
    if not node.is_running():
        raise IllegalStateException('engine not running: ' + args.engine)

    domain = node.new_domain_from_tmx(args.tmx, args.domain)
    print 'SUCCESS - imported domain "' + domain['name'] + '" with id ' + str(domain['id'])


def main():
    actions = {
        'create': main_create,
        'start': main_start,
        'stop': main_stop,
        'status': main_status,
        'delete': main_delete,
        'evaluate': main_evaluate,
        'tune': main_tune,
        'translate': main_translate,
        'add': main_add,
        '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 IllegalArgumentException as e:
        sys.stderr.write('ERROR Illegal Argument: {message}\n'.format(message=e.message))
        exit(1)
    except IllegalStateException as e:
        sys.stderr.write('ERROR Illegal State: {message}\n'.format(message=e.message))
        exit(1)
    except Exception as e:
        sys.stderr.write('ERROR Unexpected exception:\n\t{message}\n'.format(message=repr(e)))
        raise
        exit(1)


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