From f31f1511256ad523f61fa43cc25a78e221fc2b51 Mon Sep 17 00:00:00 2001
From: shekhar-rudder <85345786+shekhar-rudder@users.noreply.github.com>
Date: Thu, 11 Jul 2024 08:53:44 +0530
Subject: [PATCH] feat: add output contract in prediction model (#371)

---
 .../profiles_mlcorelib/py_native/prediction.py        | 11 +++++++++++
 1 file changed, 11 insertions(+)

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: