1 files changed, 52 insertions, 14 deletions

diff --git a/split_dataset.py b/split_dataset.py
index 9ac7bbf..a161cc1 100644
--- a/split_dataset.py
+++ b/split_dataset.py
@@ -7,10 +7,12 @@ def split_dataset(
     input_filepath: str, 
     train_filepath: str = "train.jsonl", 
     validation_filepath: str = "validation.jsonl", 
-    split_ratio: float = 0.95
+    split_ratio: float = 0.95,
+    balance_ratio: float = 1.0
 ):
     """
-    Reads a .jsonl file, shuffles it, and splits it into training and validation files.
+    Reads a .jsonl file, balances the positive/negative examples, shuffles,
+    and splits it into training and validation files.
     """
     print(f"Loading data from {input_filepath}...")
     try:
@@ -20,23 +22,51 @@ def split_dataset(
         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)
+    # 1. Separate into positive and negative examples
+    positive_examples = []
+    negative_examples = []
+    for line in lines:
+        try:
+            entry = json.loads(line)
+            if entry.get("target", "").strip():
+                positive_examples.append(line)
+            else:
+                negative_examples.append(line)
+        except json.JSONDecodeError:
+            print(f"Warning: Skipping malformed JSON line: {line.strip()}", file=sys.stderr)
 
-    # Determine the split point
-    split_index = int(len(lines) * split_ratio)
+    print(f"Found {len(positive_examples)} positive examples and {len(negative_examples)} negative examples.")
 
-    # Split the data
-    train_lines = lines[:split_index]
-    validation_lines = lines[split_index:]
+    # 2. Shuffle both lists independently
+    random.shuffle(positive_examples)
+    random.shuffle(negative_examples)
+
+    # 3. Balance the dataset based on the ratio
+    num_positives = len(positive_examples)
+    num_negatives_to_keep = int(num_positives * balance_ratio)
+
+    if len(negative_examples) < num_negatives_to_keep:
+        print(f"Warning: Not enough negative examples to meet a {balance_ratio}:1 ratio. "
+              f"Using all {len(negative_examples)} negative examples.", file=sys.stderr)
+        num_negatives_to_keep = len(negative_examples)
+
+    print(f"Balancing dataset with {num_positives} positive and {num_negatives_to_keep} negative examples.")
+    
+    balanced_lines = positive_examples + negative_examples[:num_negatives_to_keep]
+    
+    # 4. Final shuffle to mix positive and negative examples
+    print("Shuffling final balanced dataset...")
+    random.shuffle(balanced_lines)
+
+    # 5. Determine the split point and write files
+    split_index = int(len(balanced_lines) * split_ratio)
+    train_lines = balanced_lines[:split_index]
+    validation_lines = balanced_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)
@@ -44,7 +74,9 @@ def split_dataset(
     print("\nSplitting complete.")
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Split a .jsonl dataset into training and validation sets.")
+    parser = argparse.ArgumentParser(
+        description="Balance and 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", 
@@ -52,6 +84,12 @@ if __name__ == "__main__":
         default=0.95, 
         help="Ratio for the training set split (e.g., 0.95 for a 95/5 split). Default is 0.95."
     )
+    parser.add_argument(
+        "--ratio",
+        type=float,
+        default=1.0,
+        help="Ratio of negative to positive examples (e.g., 1.0 for 1:1, 2.0 for 2:1). Default is 1.0."
+    )
     
     args = parser.parse_args()
     
@@ -59,4 +97,4 @@ if __name__ == "__main__":
         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
+    split_dataset(args.input_file, split_ratio=args.split, balance_ratio=args.ratio)
\ No newline at end of file