]>
wolfpit.net Git - experiments/OpenNMT-py/.git/blob - server.py
4 from flask
import Flask
, jsonify
, request
5 from onmt
.translate
import TranslationServer
, ServerModelError
11 def start(config_file
,
12 url_root
="./translator",
16 def prefix_route(route_function
, prefix
='', mask
='{0}{1}'):
17 def newroute(route
, *args
, **kwargs
):
18 return route_function(mask
.format(prefix
, route
), *args
, **kwargs
)
22 app
.route
= prefix_route(app
.route
, url_root
)
23 translation_server
= TranslationServer()
24 translation_server
.start(config_file
)
26 @app.route('/models', methods
=['GET'])
28 out
= translation_server
.list_models()
31 @app.route('/health', methods
=['GET'])
34 out
['status'] = STATUS_OK
37 @app.route('/clone_model/<int:model_id>', methods
=['POST'])
38 def clone_model(model_id
):
40 data
= request
.get_json(force
=True)
43 timeout
= data
['timeout']
46 opt
= data
.get('opt', None)
48 model_id
, load_time
= translation_server
.clone_model(
49 model_id
, opt
, timeout
)
50 except ServerModelError
as e
:
51 out
['status'] = STATUS_ERROR
54 out
['status'] = STATUS_OK
55 out
['model_id'] = model_id
56 out
['load_time'] = load_time
60 @app.route('/unload_model/<int:model_id>', methods
=['GET'])
61 def unload_model(model_id
):
62 out
= {"model_id": model_id}
65 translation_server
.unload_model(model_id
)
66 out
['status'] = STATUS_OK
67 except Exception as e
:
68 out
['status'] = STATUS_ERROR
73 @app.route('/translate', methods
=['POST'])
75 inputs
= request
.get_json(force
=True)
78 translation
, scores
, n_best
, times
= translation_server
.run(inputs
)
79 assert len(translation
) == len(inputs
)
80 assert len(scores
) == len(inputs
)
82 out
= [[{"src": inputs
[i
]['src'], "tgt": translation
[i
],
84 "pred_score": scores
[i
]}
85 for i
in range(len(translation
))]]
86 except ServerModelError
as e
:
88 out
['status'] = STATUS_ERROR
92 @app.route('/to_cpu/<int:model_id>', methods
=['GET'])
94 out
= {'model_id': model_id}
95 translation_server
.models
[model_id
].to_cpu()
97 out
['status'] = STATUS_OK
100 @app.route('/to_gpu/<int:model_id>', methods
=['GET'])
101 def to_gpu(model_id
):
102 out
= {'model_id': model_id}
103 translation_server
.models
[model_id
].to_gpu()
105 out
['status'] = STATUS_OK
108 app
.run(debug
=debug
, host
=host
, port
=port
, use_reloader
=False,
113 parser
= configargparse
.ArgumentParser(
114 config_file_parser_class
=configargparse
.YAMLConfigFileParser
,
115 description
="OpenNMT-py REST Server")
116 parser
.add_argument("--ip", type=str, default
="0.0.0.0")
117 parser
.add_argument("--port", type=int, default
="5000")
118 parser
.add_argument("--url_root", type=str, default
="/translator")
119 parser
.add_argument("--debug", "-d", action
="store_true")
120 parser
.add_argument("--config", "-c", type=str,
121 default
="./available_models/conf.json")
125 if __name__
== '__main__':
126 parser
= _get_parser()
127 args
= parser
.parse_args()
128 start(args
.config
, url_root
=args
.url_root
, host
=args
.ip
, port
=args
.port
,