commit 4ec10c33159f435963b73da5a9bd12f25f65cda9
parent d5e69d1debc6df412f0862cb47ea90deebcd5dbf
Author: vin <git@vineetk.net>
Date: Sat, 23 Aug 2025 13:35:32 -0400
add helper script to split dataset.jsonl
Diffstat:
| A | split_dataset.py | | | 63 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
1 file changed, 63 insertions(+), 0 deletions(-)
diff --git a/split_dataset.py 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
