fix: use output_file parameter in prepare_dataset

This commit is contained in:
Christian Medina
2026-06-30 21:54:05 -04:00
parent ed0b1eddf8
commit 9cbf6549e5

View File

@@ -66,17 +66,17 @@ def prepare_dataset(input_file, output_file, test_size=0.05):
test_data = formatted[split_point:]
# Save
with open(output_dir / "train.jsonl", "w") as f:
with open(output_file / "train.jsonl", "w") as f:
for item in train_data:
f.write(json.dumps(item) + "\n")
with open(output_dir / "test.jsonl", "w") as f:
with open(output_file / "test.jsonl", "w") as f:
for item in test_data:
f.write(json.dumps(item) + "\n")
print(f"Train: {len(train_data)} examples")
print(f"Test: {len(test_data)} examples")
print(f"Saved to {output_file.parent}")
print(f"Saved to {output_file}")
def main():