Updating train-test split to Python 3: test size should be specified.

rcamino · Jan 28, 2019 · 4e4a886 · 4e4a886
1 parent 65e48ab
commit 4e4a886
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/multi_categorical_gans/datasets/README.md b/multi_categorical_gans/datasets/README.md
@@ -17,7 +17,7 @@ Example of how to split a dataset into 90% train and 10% test:
 ```bash
 python multi_categorical_gans/datasets/train_test_split.py \
     data/uscensus/USCensus1990.features.npz \
-    --percent 90 \
+    0.9 \
     data/uscensus/USCensus1990-train.features.npz \
     data/uscensus/USCensus1990-test.features.npz
 ```

diff --git a/multi_categorical_gans/datasets/train_test_split.py b/multi_categorical_gans/datasets/train_test_split.py
@@ -43,9 +43,17 @@ def main():
     features_saver = savers[options.features_format]
     features = features_loader(options.features, transform=False)
 
+    if 0 < options.train_size < 1:
+        test_size = 1 - options.train_size
+    elif options.train_size > 1:
+        test_size = len(features) - options.train_size
+    else:
+        raise Exception("Invalid train size.")
+
     if options.labels is None:
         train_features, test_features = train_test_split(features,
                                                          train_size=options.train_size,
+                                                         test_size=test_size,
                                                          shuffle=options.shuffle)
     else:
         labels_loader = loaders[options.labels_format]
@@ -55,6 +63,7 @@ def main():
         train_features, test_features, train_labels, test_labels = train_test_split(features,
                                                                                     labels,
                                                                                     train_size=options.train_size,
+                                                                                    test_size=test_size,
                                                                                     shuffle=options.shuffle)
 
     features_saver(options.train_features, train_features)