TensorFlow

import tensorflow as tf
from zenml import step, pipeline

@step
def load_dataset() -> tf.data.Dataset:
    """Step that loads and returns a tf.data.Dataset."""
    # For this example, we'll create a simple dataset
    x = tf.random.normal((100, 5))
    y = tf.random.uniform((100,), maxval=2, dtype=tf.int32)
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    return dataset.batch(32)

@step
def train_tiny_model(dataset: tf.data.Dataset) -> tf.keras.Model:
    """Step that trains a tiny model using the input dataset."""
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(10, activation='relu', input_shape=(5,)),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])

    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

    # Train for just one epoch
    model.fit(dataset, epochs=1)

    return model

@pipeline(enable_cache=False)
def tiny_model_pipeline():
    dataset = load_dataset()
    model = train_tiny_model(dataset)

tiny_model_pipeline()