<#3346 Fix tensorflow transformer> Pull request op...
# flyte-github
a
#3346 Fix tensorflow transformer Pull request opened by dansola ## Tracking issue flyteorg/flyte#6427 ## Why are the changes needed? The TensorFlow model transformer was missing the .keras file extension, causing Keras 3 to fail with "File format not supported". ## What changes were proposed in this pull request? Fixed by adding .keras extension to save/load paths and switching from MULTIPART to SINGLE blob dimensionality to match other ML transformers (PyTorch, Sklearn). ## How was this patch tested? Tested on a Union cluster here: https://demo.hosted.unionai.cloud/console/projects/flytesnacks/domains/development/executions/a6vjhxr5hsd75tmdqsfc/nodes
Copy code
import flytekit as fl
import tensorflow as tf

new_flytekit = "git+<https://github.com/flyteorg/flytekit@661edad5ea6b91d2abd4129535a80e938db01cb4>"
image = fl.ImageSpec(
    builder="union",
    apt_packages=["git"],
    packages=["union", "tensorflow", "tensorflow-datasets", "flytekitplugins-kftensorflow", new_flytekit],
)

@fl.task(container_image=image)
def train_model() -> tf.keras.Model:
    model = tf.keras.Sequential(
        [tf.keras.layers.Dense(128, activation="relu"), tf.keras.layers.Dense(10, activation="softmax")]
    )
    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return model


@fl.task(container_image=image)
def evaluate_model(model: tf.keras.Model, x: tf.Tensor, y: tf.Tensor) -> float:
    loss, accuracy = model.evaluate(x, y)
    return accuracy


@fl.workflow
def training_workflow() -> float:
    # Hardcoded tensors for testing
    x = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
    y = tf.constant([0, 1, 2, 3])

    model = train_model()
    return evaluate_model(model=model, x=x, y=y)
### Setup process ### Screenshots ## Check all the applicable boxes • I updated the documentation accordingly. • All new and existing tests passed. • All commits are signed-off. ## Related PRs ## Docs link https://bito.ai#summarystart| ### Summary by Bito • Changes the blob dimensionality from MULTIPART to SINGLE in the TensorFlow model transformer to ensure compatibility with Keras 3. • Adds the .keras file extension to save/load paths for the TensorFlow model transformer. • Updates relevant unit tests to validate the new behavior of the TensorFlow model transformer. • Overall summary: modifies TensorFlow model transformer configurations and updates unit tests to ensure compatibility with Keras 3. flyteorg/flytekit