create_dataset.py

#!/usr/bin/env python3
"""
BitTransformerLM Dataset Creation Script

Usage:
    python create_dataset.py --token YOUR_HF_TOKEN --repo-id YOUR_REPO_NAME

This script creates a comprehensive dataset for BitTransformerLM training
and uploads it to HuggingFace Hub with proper metadata and organization.
"""

import argparse
import sys
from pathlib import Path

# Add the bit_transformer module to path
sys.path.insert(0, str(Path(__file__).parent))

from bit_transformer.dataset_builder import create_bittransformerlm_dataset


def main():
    parser = argparse.ArgumentParser(description="Create BitTransformerLM Dataset")
    parser.add_argument("--token", required=True, help="HuggingFace access token")
    parser.add_argument("--repo-id", default="BitTransformerLM", help="Dataset repository ID")
    parser.add_argument("--private", action="store_true", default=True, help="Make dataset private")
    parser.add_argument("--samples", type=int, default=25000, help="Total number of samples")
    
    args = parser.parse_args()
    
    print("🚀 Starting BitTransformerLM Dataset Creation")
    print(f"Repository: {args.repo_id}")
    print(f"Private: {args.private}")
    print(f"Target samples: {args.samples}")
    print("-" * 50)
    
    try:
        dataset_url = create_bittransformerlm_dataset(
            hf_token=args.token,
            repo_id=args.repo_id
        )
        
        print("\n" + "=" * 50)
        print("🎉 SUCCESS! Dataset created and uploaded")
        print(f"📍 URL: {dataset_url}")
        print("=" * 50)
        
        print("\n📋 Next Steps:")
        print("1. View your dataset on HuggingFace Hub")
        print("2. Test loading with: `from datasets import load_dataset`")
        print("3. Integrate with BitTransformerLM training pipeline")
        print("4. Monitor dataset usage and performance metrics")
        
    except Exception as e:
        print(f"\n❌ ERROR: {e}")
        print("Please check your token and repository permissions.")
        sys.exit(1)


if __name__ == "__main__":
    main()