Update schemas for xgboost 1.7

Signed-off-by: Avi Shinnar <[email protected]>
IBM · Jan 31, 2024 · 63a5711 · 63a5711
1 parent 245d267
commit 63a5711
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/lale/lib/xgboost/xgb_classifier.py b/lale/lib/xgboost/xgb_classifier.py
@@ -898,5 +898,32 @@ def score(self, X, y):
         },
     )
 
+if xgboost_version is not None and xgboost_version >= version.Version("1.7"):
+    # https://xgboost.readthedocs.io/en/latest/python/python_api.html#module-xgboost.sklearn
+    XGBClassifier = XGBClassifier.customize_schema(
+        feature_types={
+            "description": "Used for specifying feature types without constructing a dataframe. See DMatrix for details.",
+            "laleType": "Any",
+        },
+        max_cat_threshold={
+            "description": """Maximum number of categories considered for each split.
+            Used only by partition-based splits for preventing over-fitting.
+            Also, enable_categorical needs to be set to have categorical feature support.
+            See Categorical Data and Parameters for Categorical Feature for details.""",
+            "anyOf": [
+                {
+                    "type": "integer",
+                    "minimum": 0,
+                    "distribution": "uniform",
+                    "minimumForOptimizer": 1,
+                    "maximumForOptimizer": 10,
+                },
+                {"enum": [None]},
+            ],
+            "default": None,
+        },
+        set_as_available=True,
+    )
+
 
 lale.docstrings.set_docstrings(XGBClassifier)
diff --git a/lale/lib/xgboost/xgb_regressor.py b/lale/lib/xgboost/xgb_regressor.py
@@ -833,5 +833,32 @@ def score(self, X, y):
         },
     )
 
+if xgboost_version is not None and xgboost_version >= version.Version("1.7"):
+    # https://xgboost.readthedocs.io/en/latest/python/python_api.html#module-xgboost.sklearn
+    XGBRegressor = XGBRegressor.customize_schema(
+        feature_types={
+            "description": "Used for specifying feature types without constructing a dataframe. See DMatrix for details.",
+            "laleType": "Any",
+        },
+        max_cat_threshold={
+            "description": """Maximum number of categories considered for each split.
+            Used only by partition-based splits for preventing over-fitting.
+            Also, enable_categorical needs to be set to have categorical feature support.
+            See Categorical Data and Parameters for Categorical Feature for details.""",
+            "anyOf": [
+                {
+                    "type": "integer",
+                    "minimum": 0,
+                    "distribution": "uniform",
+                    "minimumForOptimizer": 1,
+                    "maximumForOptimizer": 10,
+                },
+                {"enum": [None]},
+            ],
+            "default": None,
+        },
+        set_as_available=True,
+    )
+
 
 lale.docstrings.set_docstrings(XGBRegressor)