[PYDF] Expose data spec on the model

PiperOrigin-RevId: 576860088
google · Oct 26, 2023 · d27f7c7 · d27f7c7
1 parent 7c39e00
commit d27f7c7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/yggdrasil_decision_forests/port/python/ydf/model/BUILD b/yggdrasil_decision_forests/port/python/ydf/model/BUILD
@@ -81,6 +81,7 @@ py_library(
         ":template_cpp_export",
         # absl/logging dep,
         # numpy dep,
+        "@ydf_cc//yggdrasil_decision_forests/dataset:data_spec_py_proto",
         "@ydf_cc//yggdrasil_decision_forests/metric:metric_py_proto",
         "@ydf_cc//yggdrasil_decision_forests/model:abstract_model_py_proto",
         "//ydf/cc:ydf",

diff --git a/yggdrasil_decision_forests/port/python/ydf/model/generic_model.py b/yggdrasil_decision_forests/port/python/ydf/model/generic_model.py
@@ -21,6 +21,7 @@
 from absl import logging
 import numpy as np
 
+from yggdrasil_decision_forests.dataset import data_spec_pb2
 from yggdrasil_decision_forests.metric import metric_pb2
 from yggdrasil_decision_forests.model import abstract_model_pb2
 from ydf.cc import ydf
@@ -69,6 +70,10 @@ def describe(self, full_details: bool = False) -> str:
 
     return self._model.Describe(full_details)
 
+  def data_spec(self) -> data_spec_pb2.DataSpecification:
+    """Returns the data spec used for train the model."""
+    return self._model.data_spec()
+
   def __str__(self) -> str:
     return f"""\
 Model: {self.name()}