From 30a4df81a8e6b5e058701b8290bb7fd55f6e3b1a Mon Sep 17 00:00:00 2001
From: George Barrowclough <30464443+georyetti@users.noreply.github.com>
Date: Fri, 30 Oct 2026 15:35:42 +0000
Subject: [PATCH] fix: correct README Keras input schema example

---
 README.md | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 43e16b9..d07ffbd 100644
--- a/README.md
+++ b/README.md
@@ -45,12 +45,12 @@ fitted_pipeline = pipeline.fit(data)
 fitted_pipeline.transform(data).show()  # Use in Spark
 
 # Export for TensorFlow Serving
-tf_input_schema = [
-    {"name": "col1", "dtype": "int32", "shape": (None, 1)},
-    {"name": "col2", "dtype": "int32", "shape": (None, 1)},
-    {"name": "category", "dtype": "string", "shape": (None, 1)},
+input_schema = [
+    {"name": "col1", "dtype": "int32", "shape": (1,)},
+    {"name": "col2", "dtype": "int32", "shape": (1,)},
+    {"name": "category", "dtype": "string", "shape": (1,)},
 ]
-keras_model = fitted_pipeline.build_keras_model(tf_input_schema=tf_input_schema)
+keras_model = fitted_pipeline.build_keras_model(input_schema=input_schema)
 keras_model.save("./preprocessing_model.keras")
 ```