proper handle of lang cols

mostly-ai · Feb 19, 2025 · 69cccfa · 69cccfa
1 parent 7704136
commit 69cccfa
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/mostlyai/sdk/_local/execution/plan.py b/mostlyai/sdk/_local/execution/plan.py
@@ -14,11 +14,13 @@
 
 from pydantic import BaseModel, Field, ConfigDict
 from collections import deque
-from mostlyai.sdk.domain import ModelEncodingType, Generator, SourceTable, StepCode, TaskType
+from mostlyai.sdk.domain import ModelEncodingType, Generator, SourceTable, StepCode, TaskType, ModelType
 import uuid
 
-TABULAR_MODEL_ENCODING_TYPES = [v for v in ModelEncodingType if v.startswith("TABULAR")] + [ModelEncodingType.auto]
-LANGUAGE_MODEL_ENCODING_TYPES = [v for v in ModelEncodingType if v.startswith("LANGUAGE")]
+TABULAR_MODEL_ENCODING_TYPES = [v for v in ModelEncodingType if v.startswith(ModelType.tabular)] + [
+    ModelEncodingType.auto
+]
+LANGUAGE_MODEL_ENCODING_TYPES = [v for v in ModelEncodingType if v.startswith(ModelType.language)]
 
 TRAINING_TASK_STEPS: list[StepCode] = [
     StepCode.pull_training_data,