From d4589bd07586ba31df0a21af572b72505fe0dc97 Mon Sep 17 00:00:00 2001 From: Matt Sokoloff Date: Tue, 12 Jul 2022 08:58:34 -0400 Subject: [PATCH] support unassigned splits --- labelbox/schema/model_run.py | 8 +------- tests/integration/annotation_import/test_model_run.py | 8 +------- 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/labelbox/schema/model_run.py b/labelbox/schema/model_run.py index bc6de67f9..c698f8356 100644 --- a/labelbox/schema/model_run.py +++ b/labelbox/schema/model_run.py @@ -203,13 +203,7 @@ def assign_data_rows_to_split(self, timeout_seconds=120): split_value = split.value if isinstance(split, DataSplit) else split - - if split_value == DataSplit.UNASSIGNED.value: - raise ValueError( - f"Cannot assign split value of `{DataSplit.UNASSIGNED.value}`.") - - valid_splits = filter(lambda name: name != DataSplit.UNASSIGNED.value, - DataSplit._member_names_) + valid_splits = DataSplit._member_names_ if split_value not in valid_splits: raise ValueError( diff --git a/tests/integration/annotation_import/test_model_run.py b/tests/integration/annotation_import/test_model_run.py index 0a456f134..5042c42b7 100644 --- a/tests/integration/annotation_import/test_model_run.py +++ b/tests/integration/annotation_import/test_model_run.py @@ -146,13 +146,7 @@ def test_model_run_split_assignment(model_run, dataset, image_url): with pytest.raises(ValueError): model_run.assign_data_rows_to_split(data_row_ids, "INVALID SPLIT") - with pytest.raises(ValueError): - model_run.assign_data_rows_to_split(data_row_ids, DataSplit.UNASSIGNED) - - for split in ["TRAINING", "TEST", "VALIDATION", *DataSplit]: - if split == DataSplit.UNASSIGNED: - continue - + for split in ["TRAINING", "TEST", "VALIDATION", "UNASSIGNED", *DataSplit]: model_run.assign_data_rows_to_split(data_row_ids, split) counts = Counter() for data_row in model_run.model_run_data_rows():