add helper script to split dataset.jsonl

1 files changed, 62 insertions, 0 deletions

diff --git a/split_dataset.py b/split_dataset.py
new file mode 100644
index 0000000..9ac7bbf
--- /dev/null
+++ b/split_dataset.py
@@ -0,0 +1,62 @@
+import sys
+import json
+import random
+import argparse
+
+def split_dataset(
+    input_filepath: str, 
+    train_filepath: str = "train.jsonl", 
+    validation_filepath: str = "validation.jsonl", 
+    split_ratio: float = 0.95
+):
+    """
+    Reads a .jsonl file, shuffles it, and splits it into training and validation files.
+    """
+    print(f"Loading data from {input_filepath}...")
+    try:
+        with open(input_filepath, 'r', encoding='utf-8') as f:
+            lines = f.readlines()
+    except FileNotFoundError:
+        print(f"Error: Input file not found at {input_filepath}", file=sys.stderr)
+        sys.exit(1)
+
+    # Shuffle the dataset to ensure random distribution
+    print("Shuffling data...")
+    random.shuffle(lines)
+
+    # Determine the split point
+    split_index = int(len(lines) * split_ratio)
+
+    # Split the data
+    train_lines = lines[:split_index]
+    validation_lines = lines[split_index:]
+
+    # Write the training file
+    print(f"Writing {len(train_lines)} lines to {train_filepath}...")
+    with open(train_filepath, 'w', encoding='utf-8') as f:
+        f.writelines(train_lines)
+
+    # Write the validation file
+    print(f"Writing {len(validation_lines)} lines to {validation_filepath}...")
+    with open(validation_filepath, 'w', encoding='utf-8') as f:
+        f.writelines(validation_lines)
+    
+    print("\nSplitting complete.")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Split a .jsonl dataset into training and validation sets.")
+    parser.add_argument("input_file", type=str, help="Path to the input .jsonl file.")
+    parser.add_argument(
+        "--split", 
+        type=float, 
+        default=0.95, 
+        help="Ratio for the training set split (e.g., 0.95 for a 95/5 split). Default is 0.95."
+    )
+    
+    args = parser.parse_args()
+    
+    if not 0 < args.split < 1:
+        print("Error: Split ratio must be between 0 and 1.", file=sys.stderr)
+        sys.exit(1)
+        
+    split_dataset(args.input_file, split_ratio=args.split)
\ No newline at end of file