Deploying Tensorflow Models

from beam import App, Runtime, Image, Volume

import tensorflow as tf
import tensorflow_hub as hub


beam_cache_path = "./models"

app = App(
    name="tf-example",
    runtime=Runtime(
        cpu=1,
        memory="4Gi",
        image=Image(
            python_version="python3.8",
            python_packages=["tensorflow_hub", "tensorflow"],
        ),
    ),
    volumes=[Volume(path=beam_cache_path, name="models")],
)


@app.rest_api()
def save_model():
    model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
    # Save models to Beam Volume
    tf.saved_model.save(model, beam_cache_path)

@app.rest_api()
def load_saved_model():
    # Load the model from the Beam Volume
    loaded_model = hub.load(beam_cache_path)
    # Run the saved model
    embeddings = loaded_model(["This is an example sentence."])
    # Print model output
    print(embeddings)

beam deploy app.py:your-function-name