151 lines
4.8 KiB
Python
151 lines
4.8 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""
|
|
Convert HuggingFace OpenThoughts3-1.2M dataset to a prompt-only JSONL file
|
|
for SFT data generation (Step 1).
|
|
|
|
Extracts the prompt (user messages) from each sample and writes to JSONL.
|
|
Optionally samples a subset (default 300K) to reduce compute cost.
|
|
|
|
Usage:
|
|
python scripts/prepare_sft_prompts.py \
|
|
--output data/prompts/openthoughts3_300k.jsonl \
|
|
--num-samples 300000
|
|
|
|
# Use a local parquet file instead of downloading from HF
|
|
python scripts/prepare_sft_prompts.py \
|
|
--input data/prompts/local.parquet \
|
|
--output data/prompts/openthoughts3_300k.jsonl
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import random
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Extract prompts from OpenThoughts3-1.2M for SFT data generation."
|
|
)
|
|
parser.add_argument(
|
|
"--input", type=str, default=None,
|
|
help="Path to a local parquet/jsonl file. If not set, downloads from HuggingFace.",
|
|
)
|
|
parser.add_argument(
|
|
"--hf-dataset", type=str, default="open-thoughts/OpenThoughts3-1.2M",
|
|
help="HuggingFace dataset name (default: open-thoughts/OpenThoughts3-1.2M).",
|
|
)
|
|
parser.add_argument(
|
|
"--output", type=str, required=True,
|
|
help="Output JSONL file path.",
|
|
)
|
|
parser.add_argument(
|
|
"--num-samples", type=int, default=300000,
|
|
help="Number of samples to keep (default: 300000). Set to 0 for all.",
|
|
)
|
|
parser.add_argument(
|
|
"--seed", type=int, default=42,
|
|
help="Random seed for sampling (default: 42).",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def extract_prompt(sample):
|
|
"""Extract the prompt (non-assistant messages) from a sample.
|
|
|
|
Supports two common formats:
|
|
1. {"conversations": [{"from": "human", "value": ...}, ...]} (sharegpt)
|
|
2. {"prompt": [{"role": "user", "content": ...}, ...]} (chat messages)
|
|
"""
|
|
if "conversations" in sample:
|
|
messages = []
|
|
for turn in sample["conversations"]:
|
|
role = turn.get("from", turn.get("role", ""))
|
|
content = turn.get("value", turn.get("content", ""))
|
|
if role in ("human", "user"):
|
|
messages.append({"role": "user", "content": content})
|
|
elif role == "system":
|
|
messages.append({"role": "system", "content": content})
|
|
if messages:
|
|
return {"prompt": messages}
|
|
|
|
if "prompt" in sample:
|
|
if isinstance(sample["prompt"], list):
|
|
return {"prompt": sample["prompt"]}
|
|
elif isinstance(sample["prompt"], str):
|
|
return {"prompt": [{"role": "user", "content": sample["prompt"]}]}
|
|
|
|
if "messages" in sample:
|
|
messages = [
|
|
{"role": m["role"], "content": m["content"]}
|
|
for m in sample["messages"]
|
|
if m["role"] != "assistant"
|
|
]
|
|
if messages:
|
|
return {"prompt": messages}
|
|
|
|
return None
|
|
|
|
|
|
def load_dataset_from_hf(dataset_name):
|
|
"""Load dataset from HuggingFace."""
|
|
from datasets import load_dataset
|
|
print(f"Loading dataset from HuggingFace: {dataset_name}")
|
|
ds = load_dataset(dataset_name, split="train")
|
|
return ds
|
|
|
|
|
|
def load_dataset_from_file(path):
|
|
"""Load dataset from local file (parquet or jsonl)."""
|
|
import pandas as pd
|
|
print(f"Loading dataset from local file: {path}")
|
|
if path.endswith(".parquet"):
|
|
df = pd.read_parquet(path)
|
|
return df.to_dict("records")
|
|
elif path.endswith(".jsonl"):
|
|
with open(path) as f:
|
|
return [json.loads(line) for line in f]
|
|
else:
|
|
raise ValueError(f"Unsupported format: {path}")
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
random.seed(args.seed)
|
|
|
|
# Load dataset
|
|
if args.input:
|
|
samples = load_dataset_from_file(args.input)
|
|
else:
|
|
samples = load_dataset_from_hf(args.hf_dataset)
|
|
|
|
print(f"Total samples: {len(samples)}")
|
|
|
|
# Sample subset
|
|
if args.num_samples > 0 and args.num_samples < len(samples):
|
|
indices = random.sample(range(len(samples)), args.num_samples)
|
|
indices.sort()
|
|
samples = [samples[i] for i in indices]
|
|
print(f"Sampled {args.num_samples} samples")
|
|
|
|
# Extract prompts
|
|
from tqdm import tqdm
|
|
written = 0
|
|
skipped = 0
|
|
with open(args.output, "w") as f:
|
|
for sample in tqdm(samples, desc="Extracting prompts"):
|
|
prompt_item = extract_prompt(sample)
|
|
if prompt_item and len(prompt_item["prompt"]) > 0:
|
|
f.write(json.dumps(prompt_item) + "\n")
|
|
written += 1
|
|
else:
|
|
skipped += 1
|
|
|
|
print(f"Written: {written}, Skipped: {skipped}")
|
|
print(f"Output: {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|