def train(
self,
config,
src_file,
tgt_file,
src_vocab_info,
tgt_vocab_info,
align_file=None,
example_weights_file=None,
model_path=None,
gpuid=0,
):
# Preprocess training files.
options = config["options"].get("config", {})
options_preprocess = options.get("preprocess", {}).copy()
options_preprocess["src_vocab"] = self._convert_vocab(
config["vocabulary"]["source"]["path"])
options_preprocess["tgt_vocab"] = self._convert_vocab(
config["vocabulary"]["target"]["path"])
bin_file = os.path.join(self._data_dir, "bin")
cmd = [
"onmt_preprocess",
"-train_src",
src_file,
"-train_tgt",
tgt_file,
"-save_data",
bin_file,
]
cmd += _build_cmd(options_preprocess)
utils.run_cmd(cmd)
# Train.
options_train = options.get("train", {}).copy()
if "train_steps" not in options_train:
options_train["single_pass"] = True
options_train["train_steps"] = 0
options_train["save_checkpoint_steps"] = 0
options_train["data"] = bin_file
options_train["save_model"] = self._output_dir + "/model"
if isinstance(gpuid, list):
options_train["world_size"] = len(gpuid)
options_train["gpu_ranks"] = " ".join(str(i - 1) for i in gpuid)
elif gpuid > 0:
options_train["gpu_ranks"] = gpuid - 1
if model_path is not None:
options_train["train_from"] = os.path.join(model_path, _MODEL_NAME)
utils.run_cmd(["onmt_train"] + _build_cmd(options_train))
# Select model.
models = os.listdir(self._output_dir)
if not models:
raise RuntimeError("no model generated by the training")
if len(models) > 1:
raise RuntimeError("more than one model generated by the training")
model_file = os.path.join(self._output_dir, models[0])
return {_MODEL_NAME: model_file}
def serve(self, config, model_path, gpuid=0):
server_config_path = os.path.join(self._output_dir, "conf.json")
with open(server_config_path, "w") as server_config_file:
json.dump(
{
"models_root":
model_path,
"models": [{
"id": 0,
"model": _RELEASED_MODEL_NAME,
"opt": _trans_options(config, gpuid),
}],
},
server_config_file,
)
port = serving.pick_free_port()
process = utils.run_cmd(
[
"onmt_server",
"--ip",
"127.0.0.1",
"--port",
str(port),
"--url_root",
"/translator-backend",
"--config",
server_config_path,
],
background=True,
)
return process, {"port": port}
def serve(self, config, model_path, gpuid=0):
# Start a new tensorflow_model_server instance.
batching_parameters = self._generate_batching_parameters(config.get('serving'))
port = serving.pick_free_port()
model_name = '%s%s' % (config['source'], config['target'])
cmd = ['tensorflow_model_server',
'--port=%d' % port,
'--model_name=%s' % model_name,
'--model_base_path=%s' % model_path,
'--enable_batching=true',
'--batching_parameters_file=%s' % batching_parameters]
process = utils.run_cmd(cmd, background=True)
info = {'port': port, 'model_name': model_name}
return process, info
def serve(self, config, model_path, gpuid=0):
# Export model (deleting any previously exported models).
export_base_dir = os.path.join(model_path, "export")
if os.path.exists(export_base_dir):
shutile.rmtree(export_base_dir)
export_dir = self._export_model(config, model_path)
# Start a new tensorflow_model_server instance.
batching_parameters = self._generate_batching_parameters(
config.get('serving'))
port = serving.pick_free_port()
model_name = '%s%s' % (config['source'], config['target'])
cmd = [
'tensorflow_model_server',
'--port=%d' % port,
'--model_name=%s' % model_name,
'--model_base_path=%s' % os.path.dirname(export_dir),
'--enable_batching=true',
'--batching_parameters_file=%s' % batching_parameters
]
process = utils.run_cmd(cmd, background=True)
info = {'port': port, 'model_name': model_name}
return process, info
def _run_command(self, cmd, background=False):
return run_cmd(cmd, cwd=self._onmt_dir, background=background)
def release(self, config, model_path, gpuid=0):
model = os.path.join(model_path, _MODEL_NAME)
released_model = os.path.join(self._output_dir, _RELEASED_MODEL_NAME)
utils.run_cmd(
["onmt_release_model", "-m", model, "-o", released_model])
return {_RELEASED_MODEL_NAME: released_model}
def trans(self, config, model_path, input, output, gpuid=0):
options_trans = _trans_options(config, gpuid)
options_trans["model"] = os.path.join(model_path, _MODEL_NAME)
options_trans["src"] = input
options_trans["output"] = output
utils.run_cmd(["onmt_translate"] + _build_cmd(options_trans))
def _run_cmd(self, cmd, background=False):
return utils.run_cmd(cmd, cwd=self._onmt_py_dir, background=background)
