File size: 2,791 Bytes
58258b8 | 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 | #!/usr/bin/env python3
"""Upload SDG output JSONL to Hugging Face as a dataset.
Usage:
python -m sdg.upload_to_hf output/my_experiment/output.jsonl
python -m sdg.upload_to_hf output/my_experiment/output.jsonl --org teetone --private
"""
import argparse
import json
import tempfile
from pathlib import Path
from huggingface_hub import HfApi, create_repo
def load_jsonl(filepath: str) -> list[dict]:
"""Load JSONL file and return list of records."""
records = []
with open(filepath, "r") as f:
for line in f:
if line.strip():
records.append(json.loads(line))
return records
def main():
parser = argparse.ArgumentParser(description="Upload SDG output JSONL to Hugging Face")
parser.add_argument("jsonl_file", help="Path to the JSONL file")
parser.add_argument("--repo-name", default=None, help="Dataset repo name (default: parent directory name)")
parser.add_argument("--org", default="teetone", help="HF organization/username")
parser.add_argument("--private", action="store_true", help="Make dataset private")
args = parser.parse_args()
jsonl_path = Path(args.jsonl_file)
if not jsonl_path.exists():
raise FileNotFoundError(f"File not found: {jsonl_path}")
# Derive dataset name from parent directory (experiment_name) or filename
dataset_name = args.repo_name or jsonl_path.parent.name
repo_id = f"{args.org}/{dataset_name}"
print(f"Loading {jsonl_path}...")
records = load_jsonl(str(jsonl_path))
print(f"Loaded {len(records):,} records")
# Write clean JSONL with just messages
print("Preparing data...")
with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
temp_path = f.name
for r in records:
clean_record = {"messages": r["messages"]}
f.write(json.dumps(clean_record, ensure_ascii=False) + "\n")
api = HfApi()
print(f"Creating repo {repo_id}...")
create_repo(repo_id, repo_type="dataset", private=args.private, exist_ok=True)
print(f"Uploading to {repo_id}...")
api.upload_file(
path_or_fileobj=temp_path,
path_in_repo="data/train-00000-of-00001.jsonl",
repo_id=repo_id,
repo_type="dataset",
)
# Also upload stats.json if it exists alongside the JSONL
stats_path = jsonl_path.parent / "stats.json"
if stats_path.exists():
print(f"Uploading {stats_path}...")
api.upload_file(
path_or_fileobj=str(stats_path),
path_in_repo="stats.json",
repo_id=repo_id,
repo_type="dataset",
)
Path(temp_path).unlink()
print(f"Done! Dataset available at: https://huggingface.co/datasets/{repo_id}")
if __name__ == "__main__":
main()
|