DeepSpeech/local_tf/codegen.py
2016-12-01 17:29:34 +01:00

73 lines
2.8 KiB
Python

# depends on grpcio-tools
#
# This is intended to update the Tensorflow Serving APIs that we extract out of
# tensorflow serving repo for easier reuse. It should not be needed to run that
# except in case of changes to the .proto files with regards to the tensorflow
# serving version you are using.
#
# Running this will regenerate tensorflow/ and tensorflow_serving/. After doing
# so, please commit the changes.
from grpc.tools import protoc
import sys
import os
import shutil
import tempfile
tensorflow_root = sys.argv[1]
tensorflow_serving_root = sys.argv[2]
if not os.path.isdir(tensorflow_root) or not os.path.isdir(tensorflow_serving_root):
print "Execution requires git clones of tensorflow and tensorflow serving."
print "No tensorflow_root (%s) or no tensorflow_serving_root (%s)" % (tensorflow_root, tensorflow_serving_root)
sys.exit(1)
tmpdir = tempfile.mkdtemp(suffix='tf_proto', prefix='tmp')
thisdir = os.path.abspath('.')
distproto = os.path.join(os.path.dirname(protoc.__file__), '_proto')
tensorflow_serving_apis = os.path.join('tensorflow_serving', 'apis')
model_proto = os.path.join(tensorflow_serving_root, tensorflow_serving_apis, 'model.proto')
predict_proto = os.path.join(tensorflow_serving_root, tensorflow_serving_apis, 'predict.proto')
prediction_service = os.path.join(tensorflow_serving_root, tensorflow_serving_apis, 'prediction_service.proto')
tensorflow_core_framework = os.path.join('tensorflow', 'core', 'framework')
tensor_proto = os.path.join(tensorflow_root, tensorflow_core_framework, 'tensor.proto')
resource_handle_proto = os.path.join(tensorflow_root, tensorflow_core_framework, 'resource_handle.proto')
tensor_shape_proto = os.path.join(tensorflow_root, tensorflow_core_framework, 'tensor_shape.proto')
types_proto = os.path.join(tensorflow_root, tensorflow_core_framework, 'types.proto')
protoc.main(
(
'',
'-I%s' % (tensorflow_serving_root), '-I%s' % (os.path.join(tensorflow_serving_root, 'tensorflow')), '-I%s' % (distproto),
'--python_out=%s' % (tmpdir), '--grpc_python_out=%s' % (tmpdir),
model_proto, predict_proto, prediction_service,
)
)
protoc.main(
(
'',
'-I%s' % (tensorflow_root), '-I%s' % (distproto),
'--python_out=%s' % (tmpdir), '--grpc_python_out=%s' % (tmpdir),
tensor_proto, resource_handle_proto, tensor_shape_proto, types_proto,
)
)
def myvisit(a, dir, files):
if os.path.abspath(dir) == os.path.abspath(tmpdir):
return
initpy = os.path.join(dir, '__init__.py')
try:
os.utime(initpy, None)
except Exception:
open(initpy, 'wb').close()
os.path.walk(tmpdir, myvisit, None)
for subdir in os.listdir(tmpdir):
shutil.move(os.path.join(tmpdir, subdir), thisdir)