# isort: skip_file

# __use_pretrained_model_start__
import ray
import tensorflow as tf
from ray.train.batch_predictor import BatchPredictor
from ray.train.tensorflow import (
    to_air_checkpoint,
    TensorflowPredictor,
)


# to simulate having a pretrained model.
def build_model() -> tf.keras.Model:
    model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape=(1, )),
        tf.keras.layers.Dense(1),
    ])
    return model


model = build_model()
checkpoint = to_air_checkpoint(model)
batch_predictor = BatchPredictor(checkpoint,
                                 TensorflowPredictor,
                                 model_definition=build_model)
predict_dataset = ray.data.range(3)
predictions = batch_predictor.predict(predict_dataset)

# __use_pretrained_model_end__
Example #2
0
# __batch_prediction_start__
import pandas as pd
from ray.train.batch_predictor import BatchPredictor

batch_predictor = BatchPredictor(checkpoint,
                                 TensorflowPredictor,
                                 model_definition=build_model)
# Create a dummy dataset.
ds = ray.data.from_pandas(
    pd.DataFrame({
        "feature_1": [1, 2, 3],
        "label": [1, 2, 3]
    }))

# Use `feature_columns` to specify the input columns to your model.
predictions = batch_predictor.predict(ds, feature_columns=["feature_1"])
print(predictions.show())

# {'predictions': array([-1.2789773], dtype=float32)}
# {'predictions': array([-2.5579545], dtype=float32)}
# {'predictions': array([-3.8369317], dtype=float32)}
# __batch_prediction_end__


# __compute_accuracy_start__
def calculate_accuracy(df):
    return pd.DataFrame({"correct": int(df["predictions"][0]) == df["label"]})


predictions = batch_predictor.predict(ds,
                                      feature_columns=["feature_1"],