Add t2e models to Enum

Author: kalama-ai

octopus/types.py

@@ -45,7 +45,7 @@ class ImputationMethod(str, Enum):
 
 
 class ModelName(StrEnum):
-    """Available model names for user-friendly model selection with IDE autocomplete.
+    """Available model names.
 
     Use this enum for IDE autocomplete when specifying models, e.g.::
 
@@ -80,6 +80,10 @@ class ModelName(StrEnum):
     HistGradientBoostingRegressor = "HistGradientBoostingRegressor"
     TabularNNRegressor = "TabularNNRegressor"
 
+    # Time-to-event (survival) models
+    CatBoostCoxSurvival = "CatBoostCoxSurvival"
+    XGBoostCoxSurvival = "XGBoostCoxSurvival"
+
 
 class FIResultLabel(StrEnum):
     """Labels used in feature-importance result DataFrames.