acoustic-carpenter-78188
11/01/2025, 1:41 AMimport 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/flytekitacoustic-carpenter-78188
11/01/2025, 5:17 PM