1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
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,
balance_ratio: float = 1.0
):
"""
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:
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)
# 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)
print(f"Found {len(positive_examples)} positive examples and {len(negative_examples)} negative examples.")
# 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:]
print(f"Writing {len(train_lines)} lines to {train_filepath}...")
with open(train_filepath, 'w', encoding='utf-8') as f:
f.writelines(train_lines)
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="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",
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."
)
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()
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, balance_ratio=args.ratio)
|
