fix tests

drivendataorg · Sep 27, 2022 · 0fdbc20 · 0fdbc20
1 parent c8b39be
commit 0fdbc20
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -117,7 +117,11 @@ def test_extra_column(tmp_path, labels_absolute_path):
         index=False,
     )
     # this column is not one hot encoded
-    config = TrainConfig(labels=tmp_path / "extra_species_col.csv", save_dir=tmp_path / "my_model")
+    config = TrainConfig(
+        labels=tmp_path / "extra_species_col.csv",
+        save_dir=tmp_path / "my_model",
+        use_default_model_labels=False,
+    )
     assert list(config.labels.columns) == [
         "filepath",
         "split",

diff --git a/zamba/models/config.py b/zamba/models/config.py
@@ -633,7 +633,7 @@ def make_split(labels, values):
         too_few = {
             k.split("species_", 1)[1]: v
             for k, v in num_videos_per_species.items()
-            if v < len(expected_splits)
+            if 0 < v < len(expected_splits)
         }
 
         if len(too_few) > 0:
@@ -644,12 +644,14 @@ def make_split(labels, values):
         for c in labels.filter(regex="species_").columns:
             species_df = labels[labels[c] > 0]
 
-            # within each species, seed splits by putting one video in each set and then allocate videos based on split proportions
-            labels.loc[species_df.index, "split"] = expected_splits + random.choices(
-                list(values["split_proportions"].keys()),
-                weights=list(values["split_proportions"].values()),
-                k=len(species_df) - len(expected_splits),
-            )
+            if len(species_df):
+
+                # within each species, seed splits by putting one video in each set and then allocate videos based on split proportions
+                labels.loc[species_df.index, "split"] = expected_splits + random.choices(
+                    list(values["split_proportions"].keys()),
+                    weights=list(values["split_proportions"].values()),
+                    k=len(species_df) - len(expected_splits),
+                )
 
         logger.info(f"{labels.split.value_counts()}")