diff --git a/src/predictions/profiles_mlcorelib/py_native/prediction.py b/src/predictions/profiles_mlcorelib/py_native/prediction.py
index f2ab9aed0..cb5fdee9f 100644
--- a/src/predictions/profiles_mlcorelib/py_native/prediction.py
+++ b/src/predictions/profiles_mlcorelib/py_native/prediction.py
@@ -64,6 +64,17 @@ class PredictionModel(BaseModelType):
     }
 
     def __init__(self, build_spec: dict, schema_version: int, pb_version: str) -> None:
+        build_spec["contract"] = {
+            "with_entity_ids": [build_spec["entity_key"]],
+            "with_columns": [
+                {
+                    "name": build_spec["ml_config"]["outputs"]["column_names"][
+                        "percentile"
+                    ]
+                },
+                {"name": build_spec["ml_config"]["outputs"]["column_names"]["score"]},
+            ],
+        }
         super().__init__(build_spec, schema_version, pb_version)
 
     def get_material_recipe(self) -> PyNativeRecipe: