[tab] gdr/generate_subset.py (python)
from datasets import load_dataset, Dataset, DatasetDict, concatenate_datasets
from huggingface_hub import login, create_repo
import os
from dotenv import load_dotenv
import random
from tqdm import tqdm
import gc
import sys

# Load environment variables from .env file
load_dotenv()

# Load configuration from environment variables
HF_USER = os.getenv("HF_USER")
NEW_DS_NAME = os.getenv("NEW_DS_NAME")
HF_TOKEN = os.getenv("HF_TOKEN")

SOURCE_DS = "nvidia/Llama-Nemotron-Post-Training-Dataset"
CONFIG = "SFT"

target_counts = {
    "math": 100_000,
    "code": 50_000,
    "science": 10_000
}

GENERATOR_FILTER = None
SEED = 42
PRIVATE = False

def sample_split(split_name, num_samples):
    print(f"
Sampling {num_samples} examples from split '{split_name}'")

    ds_stream = load_dataset(SOURCE_DS, CONFIG, split=split_name, streaming=True)

    def filter_examples(example):
        if GENERATOR_FILTER is not None and example.get("generator") != GENERATOR_FILTER:
            return False
        return example.get("reasoning") == "on"

    filtered_ds = ds_stream.filter(filter_examples)
    
    sampled = []
    count = 0
    reservoir_size = num_samples

    # Progress bar for sampling
    pbar = tqdm(desc=f"Sampling {split_name}", unit="examples", dynamic_ncols=True)
    
    for i, example in enumerate(filtered_ds, start=1):
        if count < reservoir_size:
            sampled.append(example)
            count += 1
            pbar.update(1)
            pbar.set_postfix(collected=count, target=reservoir_size)
        else:
            # Option: Stop early once we have enough samples
            # Comment out the next 4 lines if you want true reservoir sampling
            break
            
            # Reservoir sampling (keeps going to maintain randomness)
            j = random.randint(1, i)
            if j <= reservoir_size:
                sampled[j - 1] = example
        
        # Update progress bar every 1000 iterations when reservoir is full
        if count >= reservoir_size and i % 1000 == 0:
            pbar.set_postfix(collected=count, processed=i, target=reservoir_size)

    pbar.close()
    print(f" → Collected {len(sampled)} examples from split '{split_name}'")
    return sampled

def main():
    random.seed(SEED)
    
    # Validate required environment variables
    if not HF_TOKEN:
        raise ValueError("HF_TOKEN environment variable is required")
    if not HF_USER:
        raise ValueError("HF_USER environment variable is required")
    if not NEW_DS_NAME:
        raise ValueError("NEW_DS_NAME environment variable is required")
    
    login(token=HF_TOKEN)

    print(f"
{'='*50}")
    print(f"Starting dataset sampling process...")
    print(f"Target counts: {target_counts}")
    print(f"{'='*50}")

    split_samples = {}
    for split_name, n in tqdm(target_counts.items(), desc="Processing splits", unit="split"):
        sampled = sample_split(split_name, n)
        split_samples[split_name] = sampled

    sample_feature_names = None
    for split_name, examples in split_samples.items():
        if examples:
            sample_feature_names = list(examples[0].keys())
            break
    
    if sample_feature_names is None:
        raise ValueError("No examples were collected. Check your filters or target_counts.")

    print(f"
{'='*50}")
    print("Combining and processing samples...")
    print(f"{'='*50}")
    
    combined_list = []
    total_examples = sum(len(examples) for examples in split_samples.values())
    
    with tqdm(total=total_examples, desc="Combining samples", unit="examples") as pbar:
        for split_name, examples in split_samples.items():
            for ex in examples:
                ex_copy = dict(ex)
                ex_copy["split_from"] = split_name
                combined_list.append(ex_copy)
                pbar.update(1)

    print(f"
Total combined examples = {len(combined_list)}")
    print("Shuffling dataset...")
    
    random.shuffle(combined_list)

    print("Creating dataset structure...")
    print(f"Processing {len(combined_list)} examples...")
    
    # Process in batches to avoid memory issues
    batch_size = 10000
    if len(combined_list) > batch_size:
        print(f"Large dataset detected. Processing in batches of {batch_size}...")
        
        # Create dataset in batches
        batches = []
        for i in tqdm(range(0, len(combined_list), batch_size), desc="Creating batches", unit="batch"):
            batch = combined_list[i:i + batch_size]
            batch_ds = Dataset.from_list(batch)
            batches.append(batch_ds)
        
        print("Concatenating batches...")
        final_ds = concatenate_datasets(batches)
        
        # Clear memory
        del batches
        del combined_list
        gc.collect()
    else:
        final_ds = Dataset.from_list(combined_list)
    
    ds_to_upload = DatasetDict({"subset": final_ds})

    repo_id = f"{HF_USER}/{NEW_DS_NAME}"
    print(f"
Creating repository '{repo_id}'...")
    create_repo(repo_id, repo_type="dataset", private=PRIVATE, exist_ok=True)
    print(f"Repository '{repo_id}' ready.")

    print(f"
Uploading dataset to Hugging Face Hub...")
    print("(Note: Upload progress depends on HF Hub - this may take a while)")
    ds_to_upload.push_to_hub(repo_id, private=PRIVATE)
    print(f"
🎉 Upload complete! Dataset available at: https://huggingface.co/datasets/{repo_id}")

if __name__ == "__main__":
    main()