Translation server and demo frontend

examples/server.json

@@ -0,0 +1,14 @@
+{
+    "models_root": "./models",
+    "models": [
+        {
+            "id": 0,
+            "opts": {
+                "config": "out.modernize.cpu.yaml",
+                "model": "models/smoketest.modernize_step_10",
+                "task_id": "train_de-fr_split0"
+            },
+            "models": ["smoketest.modernize_step_10"]
+        }
+    ]
+}

mammoth/bin/server.py

@@ -78,10 +78,16 @@ def unload_model(model_id):
 
     @app.route('/translate', methods=['POST'])
     def translate():
-        inputs = request.get_json(force=True)
+        out = {}
+        inputs = request.get_json(force=True, silent=True)
+        if inputs is None:
+            error = f'Unable to parse {request.data}'
+            logger.warning(error)
+            out['error'] = str(error)
+            out['status'] = STATUS_ERROR
+            return jsonify(out)
         if debug:
             logger.info(inputs)
-        out = {}
         try:
             trans, scores, n_best, _, aligns = translation_server.run(inputs)
             assert len(trans) == len(inputs) * n_best

mammoth/inputters/dataset.py

@@ -3,6 +3,7 @@
 import itertools
 from functools import partial
 import gzip
+from io import IOBase
 
 import torch
 from torch.nn.utils.rnn import pad_sequence
@@ -54,12 +55,16 @@ def _make_example_dict(packed):
             'line_idx': next(line_idx_generator)
         }
 
-    if src_path.endswith('.gz'):
+    if isinstance(src_path, IOBase):
+        src_fh = src_path
+    elif src_path.endswith('.gz'):
         src_fh = gzip.open(src_path, 'rt')
     else:
         src_fh = open(src_path, 'rt')
     if tgt_path is None:
         tgt_fh = itertools.repeat(None)
+    elif isinstance(tgt_path, IOBase):
+        tgt_fh = src_path
     elif tgt_path.endswith('.gz'):
         tgt_fh = gzip.open(tgt_path, 'rt')
     else:

mammoth/tests/test_data_prepare.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-# # # FIXME: Broken in FoTraNMT
+# # # FIXME: Broken in MAMMOTH
 #
 # import copy
 # import unittest

mammoth/tests/test_models.py

@@ -141,7 +141,7 @@ def test_method(self):
 '''
 opts.brnn = False
 
-# FIXME: Most tests disabled: FoTraNMT only supports Transformer
+# FIXME: Most tests disabled: MAMMOTH only supports Transformer
 test_embeddings = [
     # [],
     [('decoder_type', 'transformer')]
@@ -150,7 +150,7 @@ def test_method(self):
 for p in test_embeddings:
     _add_test(p, 'embeddings_forward')
 
-# FIXME: All tests disabled: FoTraNMT only supports Transformer, and the test for Transformer is broken
+# FIXME: All tests disabled: MAMMOTH only supports Transformer, and the test for Transformer is broken
 tests_encoder = [
     # [],
     # [('encoder_type', 'mean')],
@@ -161,7 +161,7 @@ def test_method(self):
 for p in tests_encoder:
     _add_test(p, 'encoder_forward')
 
-# FIXME: Most tests disabled: FoTraNMT only supports Transformer
+# FIXME: Most tests disabled: MAMMOTH only supports Transformer
 tests_nmtmodel = [
     # [('rnn_type', 'GRU')],
     # [('layers', 10)],
@@ -193,6 +193,6 @@ def test_method(self):
 ]
 
 
-# ## FIXME: Broken in FoTraNMT
+# ## FIXME: Broken in MAMMOTH
 # for p in tests_nmtmodel:
 #     _add_test(p, 'nmtmodel_forward')

mammoth/tests/test_text_dataset.py

@@ -1,4 +1,4 @@
-# FIXME, or rather TODO in FoTran
+# FIXME, or rather TODO in MAMMOTH
 # import unittest
 #
 # import itertools

mammoth/tests/test_translation_server.py

@@ -13,7 +13,7 @@
 
 
 class TestServerModel(unittest.TestCase):
-    @unittest.skip('Broken in FoTraNMT')  # FIXME
+    @unittest.skip('Broken in MAMMOTH')  # FIXME
     def test_deferred_loading_model_and_unload(self):
         model_id = 0
         opts = {"models": ["test_model.pt"]}
@@ -26,7 +26,7 @@ def test_deferred_loading_model_and_unload(self):
         sm.unload()
         self.assertFalse(sm.loaded)
 
-    @unittest.skip('Broken in FoTraNMT')  # FIXME
+    @unittest.skip('Broken in MAMMOTH')  # FIXME
     def test_load_model_on_init_and_unload(self):
         model_id = 0
         opts = {"models": ["test_model.pt"]}
@@ -37,7 +37,7 @@ def test_load_model_on_init_and_unload(self):
         sm.unload()
         self.assertFalse(sm.loaded)
 
-    @unittest.skip('Broken in FoTraNMT')  # FIXME
+    @unittest.skip('Broken in MAMMOTH')  # FIXME
     def test_tokenizing_with_no_tokenizer_fails(self):
         model_id = 0
         opts = {"models": ["test_model.pt"]}
@@ -46,7 +46,7 @@ def test_tokenizing_with_no_tokenizer_fails(self):
         with self.assertRaises(ValueError):
             sm.tokenize("hello world")
 
-    @unittest.skip('Broken in FoTraNMT')  # FIXME
+    @unittest.skip('Broken in MAMMOTH')  # FIXME
     def test_detokenizing_with_no_tokenizer_fails(self):
         model_id = 0
         opts = {"models": ["test_model.pt"]}
@@ -57,7 +57,7 @@ def test_detokenizing_with_no_tokenizer_fails(self):
 
     if torch.cuda.is_available():
 
-        @unittest.skip('Broken in FoTraNMT')  # FIXME
+        @unittest.skip('Broken in MAMMOTH')  # FIXME
         def test_moving_to_gpu_and_back(self):
             torch.cuda.set_device(torch.device("cuda", 0))
             model_id = 0
@@ -74,7 +74,7 @@ def test_moving_to_gpu_and_back(self):
             for p in sm.translator.model.parameters():
                 self.assertEqual(p.device.type, "cpu")
 
-        @unittest.skip('Broken in FoTraNMT')  # FIXME
+        @unittest.skip('Broken in MAMMOTH')  # FIXME
         def test_initialize_on_gpu_and_move_back(self):
             torch.cuda.set_device(torch.device("cuda", 0))
             model_id = 0
@@ -111,7 +111,7 @@ def test_initialize_on_nonzero_gpu_and_back(self):
                 for p in sm.translator.model.parameters():
                     self.assertEqual(p.device.type, "cpu")
 
-    @unittest.skip('Broken in FoTraNMT')  # FIXME
+    @unittest.skip('Broken in MAMMOTH')  # FIXME
     def test_run(self):
         model_id = 0
         opts = {"models": ["test_model.pt"]}

mammoth/translate/translation_server.py

@@ -12,16 +12,19 @@
 import torch
 import mammoth.opts
 
+from io import StringIO
 from itertools import islice, zip_longest
 from copy import deepcopy
 
 from mammoth.constants import DefaultTokens
+from mammoth.distributed import TaskSpecs
+from mammoth.distributed.tasks import get_adapter_ids
+from mammoth.translate.translator import build_translator
+from mammoth.transforms import get_transforms_cls, make_transforms, TransformPipe
+from mammoth.utils.alignment import to_word_align
 from mammoth.utils.logging import init_logger
 from mammoth.utils.misc import set_random_seed
-from mammoth.utils.misc import check_model_config
-from mammoth.utils.alignment import to_word_align
 from mammoth.utils.parse import ArgumentParser
-from mammoth.translate.translator import build_translator
 
 
 def critical(func):
@@ -175,7 +178,7 @@ def start(self, config_file):
                                         parameter for model #%d"""
                         % i
                     )
-            check_model_config(conf, self.models_root)
+            # check_model_config(conf, self.models_root)
             kwargs = {
                 'timeout': conf.get('timeout', None),
                 'load': conf.get('load', None),
@@ -383,7 +386,11 @@ def parse_opt(self, opts):
         prec_argv = sys.argv
         sys.argv = sys.argv[:1]
         parser = ArgumentParser()
-        mammoth.opts.translate_opts(parser)
+
+        parser.translation = True
+
+        mammoth.opts.dynamic_prepare_opts(parser)
+        mammoth.opts.translate_opts(parser, dynamic=True)
 
         models = opts['models']
         if not isinstance(models, (list, tuple)):
@@ -397,11 +404,19 @@ def parse_opt(self, opts):
                 sys.argv += [str(model) for model in v]
             elif isinstance(v, bool):
                 sys.argv += ['-%s' % k]
+            elif k == 'transforms':
+                if type(v) is str:
+                    sys.argv += ['-transforms', v]
+                else:
+                    assert type(v) is list
+                    sys.argv += ['-transforms', *v]
             else:
                 sys.argv += ['-%s' % k, str(v)]
 
         opts = parser.parse_args()
+        ArgumentParser.validate_prepare_opts(opts)
         ArgumentParser.validate_translate_opts(opts)
+        ArgumentParser.validate_translate_opts_dynamic(opts)
         opts.cuda = opts.gpu > -1
 
         sys.argv = prec_argv
@@ -431,9 +446,27 @@ def load(self, preload=False):
                     preload=preload,
                 )
             else:
+                task = self.make_task(self.opts)
                 self.translator = build_translator(
-                    self.opts, report_score=False, out_file=codecs.open(os.devnull, "w", "utf-8")
+                    self.opts,
+                    task,
+                    report_score=False,
+                    out_file=codecs.open(os.devnull, "w", "utf-8"),
                 )
+
+                # Build transforms
+                transforms_cls = get_transforms_cls(self.opts._all_transform)
+                transforms = make_transforms(
+                    self.opts,
+                    transforms_cls,
+                    self.translator.vocabs,
+                    task=task
+                )
+                data_transform = [
+                    transforms[name] for name in self.opts.transforms if name in transforms
+                ]
+
+                self.transform = TransformPipe.build_from(data_transform)
         except RuntimeError as e:
             raise ServerModelError("Runtime Error: %s" % str(e))
 
@@ -442,6 +475,36 @@ def load(self, preload=False):
         self.reset_unload_timer()
         self.loading_lock.set()
 
+    def make_task(self, opts):
+        corpus_id = opts.task_id
+        print(f'opts.tasks {type(opts.tasks)}')
+        corpus_opts = opts.tasks[corpus_id]
+        src_lang, tgt_lang = corpus_opts['src_tgt'].split('-', 1)
+        encoder_id = corpus_opts.get('enc_sharing_group', [src_lang])
+        decoder_id = corpus_opts.get('dec_sharing_group', [tgt_lang])
+        if 'adapters' in corpus_opts:
+            encoder_adapter_ids = get_adapter_ids(opts, corpus_opts, 'encoder')
+            decoder_adapter_ids = get_adapter_ids(opts, corpus_opts, 'decoder')
+        else:
+            encoder_adapter_ids = None
+            decoder_adapter_ids = None
+        task = TaskSpecs(
+            node_rank=None,
+            local_rank=None,
+            src_lang=src_lang,
+            tgt_lang=tgt_lang,
+            encoder_id=encoder_id,
+            decoder_id=decoder_id,
+            corpus_id=corpus_id,
+            weight=1.0,
+            corpus_opts=corpus_opts,
+            src_vocab=None,
+            tgt_vocab=None,
+            encoder_adapter_ids=encoder_adapter_ids,
+            decoder_adapter_ids=decoder_adapter_ids,
+        )
+        return task
+
     @critical
     def run(self, inputs):
         """Translate `inputs` using this model
@@ -514,8 +577,9 @@ def run(self, inputs):
 
         if len(texts_to_translate) > 0:
             try:
-                scores, predictions = self.translator.translate(
-                    texts_to_translate,
+                scores, predictions = self.translator.translate_dynamic(
+                    src=StringIO("\n".join(texts_to_translate)),
+                    transform=self.transform,
                     tgt=texts_ref,
                     batch_size=len(texts_to_translate) if self.opts.batch_size == 0 else self.opts.batch_size,
                 )

mammoth/translate/translator.py

@@ -449,11 +449,11 @@ def _translate(
                 of `n_best` predictions
         """
 
-        self.logger.info("src vocab: {}".format(self.vocabs['src']))
-        self.logger.info("transforms: {}".format(transforms))
+        self._log("src vocab: {}".format(self.vocabs['src']))
+        self._log("transforms: {}".format(transforms))
         corpus = ParallelCorpus(
-            self.src_file_path,
-            self.tgt_file_path,  # may be None
+            src,
+            tgt,
             self.vocabs['src'],
             self.vocabs['tgt'],
             transforms=transforms,  # I suppose you might want *some* transforms

mammoth/utils/misc.py

@@ -5,6 +5,7 @@
 import inspect
 import numpy as np
 from itertools import islice, repeat
+from io import StringIO
 import os
 
 
@@ -30,16 +31,16 @@ def split_corpus(path, shard_size, default=None):
 
 
 def _split_corpus(path, shard_size):
-    """Yield a `list` containing `shard_size` line of `path`."""
-    with open(path, "rb") as f:
+    """Yield io's with `shard_size` lines each."""
+    with open(path, "rt") as f:
         if shard_size <= 0:
-            yield f.readlines()
+            yield f
         else:
             while True:
-                shard = list(islice(f, shard_size))
+                shard = "\n".join(islice(f, shard_size))
                 if not shard:
                     break
-                yield shard
+                yield StringIO(shard)
 
 
 def aeq(*args):

tools/demo/README.md

@@ -0,0 +1,9 @@
+Run the backend (translation server)
+
+    python mammoth/bin/server.py --config server.json
+
+Run the frontend
+
+    streamlit run mammoth_demo.py
+
+Start the backend first, otherwise the frontend will show an error until refreshed.