split_dataset.py (3688B)
1 import sys 2 import json 3 import random 4 import argparse 5 6 def split_dataset( 7 input_filepath: str, 8 train_filepath: str = "train.jsonl", 9 validation_filepath: str = "validation.jsonl", 10 split_ratio: float = 0.95, 11 balance_ratio: float = 1.0 12 ): 13 """ 14 Reads a .jsonl file, balances the positive/negative examples, shuffles, 15 and splits it into training and validation files. 16 """ 17 print(f"Loading data from {input_filepath}...") 18 try: 19 with open(input_filepath, 'r', encoding='utf-8') as f: 20 lines = f.readlines() 21 except FileNotFoundError: 22 print(f"Error: Input file not found at {input_filepath}", file=sys.stderr) 23 sys.exit(1) 24 25 # 1. Separate into positive and negative examples 26 positive_examples = [] 27 negative_examples = [] 28 for line in lines: 29 try: 30 entry = json.loads(line) 31 if entry.get("target", "").strip(): 32 positive_examples.append(line) 33 else: 34 negative_examples.append(line) 35 except json.JSONDecodeError: 36 print(f"Warning: Skipping malformed JSON line: {line.strip()}", file=sys.stderr) 37 38 print(f"Found {len(positive_examples)} positive examples and {len(negative_examples)} negative examples.") 39 40 # 2. Shuffle both lists independently 41 random.shuffle(positive_examples) 42 random.shuffle(negative_examples) 43 44 # 3. Balance the dataset based on the ratio 45 num_positives = len(positive_examples) 46 num_negatives_to_keep = int(num_positives * balance_ratio) 47 48 if len(negative_examples) < num_negatives_to_keep: 49 print(f"Warning: Not enough negative examples to meet a {balance_ratio}:1 ratio. " 50 f"Using all {len(negative_examples)} negative examples.", file=sys.stderr) 51 num_negatives_to_keep = len(negative_examples) 52 53 print(f"Balancing dataset with {num_positives} positive and {num_negatives_to_keep} negative examples.") 54 55 balanced_lines = positive_examples + negative_examples[:num_negatives_to_keep] 56 57 # 4. Final shuffle to mix positive and negative examples 58 print("Shuffling final balanced dataset...") 59 random.shuffle(balanced_lines) 60 61 # 5. Determine the split point and write files 62 split_index = int(len(balanced_lines) * split_ratio) 63 train_lines = balanced_lines[:split_index] 64 validation_lines = balanced_lines[split_index:] 65 66 print(f"Writing {len(train_lines)} lines to {train_filepath}...") 67 with open(train_filepath, 'w', encoding='utf-8') as f: 68 f.writelines(train_lines) 69 70 print(f"Writing {len(validation_lines)} lines to {validation_filepath}...") 71 with open(validation_filepath, 'w', encoding='utf-8') as f: 72 f.writelines(validation_lines) 73 74 print("\nSplitting complete.") 75 76 if __name__ == "__main__": 77 parser = argparse.ArgumentParser( 78 description="Balance and split a .jsonl dataset into training and validation sets." 79 ) 80 parser.add_argument("input_file", type=str, help="Path to the input .jsonl file.") 81 parser.add_argument( 82 "--split", 83 type=float, 84 default=0.95, 85 help="Ratio for the training set split (e.g., 0.95 for a 95/5 split). Default is 0.95." 86 ) 87 parser.add_argument( 88 "--ratio", 89 type=float, 90 default=1.0, 91 help="Ratio of negative to positive examples (e.g., 1.0 for 1:1, 2.0 for 2:1). Default is 1.0." 92 ) 93 94 args = parser.parse_args() 95 96 if not 0 < args.split < 1: 97 print("Error: Split ratio must be between 0 and 1.", file=sys.stderr) 98 sys.exit(1) 99 100 split_dataset(args.input_file, split_ratio=args.split, balance_ratio=args.ratio)