TensorFlow Hub + Cloud Run

app.py

import os

from flask import Flask, request, jsonify

import tensorflow_hub as hub
import tensorflow_text


app = Flask(__name__)


model = "https://tfhub.dev/google/nnlm-ja-dim50-with-normalization/2"
module = hub.load(model)

@app.route('/')
def get_embedding():
    sentence = request.args.get('sentence', '')
    emb = module([sentence])
    return jsonify({'embedding': emb.numpy().tolist()})

if __name__ == "__main__":
    app.run(host='0.0.0.0',port=int(os.environ.get('PORT', 8080)))

Dockerfile

FROM python:3.7

ENV APP_HOME /app
WORKDIR $APP_HOME
COPY . .

RUN pip install Flask gunicorn tensorflow tensorflow-hub tensorflow_text>=2.0.0rc0
RUN python load_module.py

CMD exec gunicorn --bind :$PORT --workers 1 --threads 8 app:app

load_module.py

import tensorflow_hub as hub
import tensorflow_text

model = "https://tfhub.dev/google/nnlm-ja-dim50-with-normalization/2"

if __name__ == '__main__':
    hub.load(model)