Improvement in train_test_split function

dshahid380 · dshahid380 · commit 1c3a5ca5db24 · 2019-04-12T21:05:57.000+05:30
diff --git a/learning.py b/learning.py
@@ -1049,14 +1049,35 @@ def grade_learner(predict, tests):
     return mean(int(predict(X) == y) for X, y in tests)
 
 
-def train_test_split(dataset, start, end):
-    """Reserve dataset.examples[start:end] for test; train on the remainder."""
-    start = int(start)
-    end = int(end)
-    examples = dataset.examples
-    train = examples[:start] + examples[end:]
-    val = examples[start:end]
-    return train, val
+def train_test_split(dataset, start = None, end = None, test_split = None):
+    """If you are giving 'start' and 'end' as a parameter,
+    then it will return testing set from index 'start' to 'end'
+    and rest for training.
+    If you give 'test_split' as parameter then it will first shuffle the
+    dataset then return test_split * 100% as testing set and rest as 
+    training set.
+    """
+    
+    if start == None and end != None:
+        raise ValueError("'start' parameter is missing")
+
+    if start != None and end == None:
+        raise ValueError("'end' parameter is missing")
+
+    if test_split == None:
+        examples = dataset.examples
+        train = examples[:start] + examples[end:]
+        val = examples[start:end]
+        return train, val
+    else:
+        examples = dataset.examples
+        total_size = len(examples)
+        val_size = int(total_size * test_split)
+        train_size = total_size - val_size
+        random.shuffle(examples)
+        train = examples[:train_size]
+        val = examples[train_size:total_size]
+        return train, val
 
 
 def cross_validation(learner, size, dataset, k=10, trials=1):
