dataset_compressor.py

#!/usr/bin/env python3
"""
HuggingFace Dataset Compressor
Converts any HuggingFace conversational dataset to compressed UTF-8 format
"""

import json
import subprocess
import sys
from pathlib import Path
from datasets import load_dataset
import argparse
from tqdm import tqdm

def compress_text(text):
    """Use C# compression to compress text"""
    try:
        result = subprocess.run([
            'dotnet', 'run', '--project', 'UTF8TranslationLayer', 
            'compress', text
        ], capture_output=True, text=True, cwd='.')
        
        if result.returncode == 0:
            return result.stdout.strip()
        else:
            print(f"Compression error: {result.stderr}")
            return None
    except Exception as e:
        print(f"Compression failed: {e}")
        return None

def extract_conversation_pairs(example, dataset_format):
    """Extract user/assistant pairs from different dataset formats"""
    pairs = []
    
    if dataset_format == 'ultrachat':
        # ultrachat_200k format: messages list with role/content
        if 'messages' in example:
            messages = example['messages']
            for i in range(0, len(messages)-1, 2):
                if (i+1 < len(messages) and 
                    messages[i].get('role') == 'user' and 
                    messages[i+1].get('role') == 'assistant'):
                    pairs.append({
                        'user': messages[i]['content'],
                        'assistant': messages[i+1]['content']
                    })
    
    elif dataset_format == 'alpaca':
        # Alpaca format: instruction/output
        if 'instruction' in example and 'output' in example:
            pairs.append({
                'user': example['instruction'],
                'assistant': example['output']
            })
    
    elif dataset_format == 'sharegpt':
        # ShareGPT format: conversations list
        if 'conversations' in example:
            convs = example['conversations']
            for i in range(0, len(convs)-1, 2):
                if (i+1 < len(convs) and 
                    convs[i].get('from') == 'human' and 
                    convs[i+1].get('from') == 'gpt'):
                    pairs.append({
                        'user': convs[i]['value'],
                        'assistant': convs[i+1]['value']
                    })
    
    return pairs

def process_dataset(dataset_name, dataset_format, output_file, max_examples=None, split='train'):
    """Download and compress a HuggingFace dataset"""
    print(f"📦 Loading dataset: {dataset_name}")
    
    try:
        # Stream the dataset instead of loading all at once
        dataset = load_dataset(dataset_name, split=split, streaming=True)
        print(f"✅ Dataset streaming enabled")
        
        # If max_examples specified, take only that many
        if max_examples:
            dataset = dataset.take(max_examples)
            print(f"🔢 Limited to {max_examples} examples")
        
    except Exception as e:
        print(f"❌ Failed to load dataset: {e}")
        return
    
    compressed_data = []
    successful_compressions = 0
    
    print(f"🗜️ Compressing conversations...")
    
    for i, example in enumerate(tqdm(dataset)):
        pairs = extract_conversation_pairs(example, dataset_format)
        
        for pair in pairs:
            user_text = pair['user']
            assistant_text = pair['assistant']
            
            # Skip very long texts to avoid compression issues
            if len(user_text) > 1000 or len(assistant_text) > 1000:
                continue
            
            # Compress both user and assistant text
            compressed_user = compress_text(user_text)
            compressed_assistant = compress_text(assistant_text)
            
            if compressed_user and compressed_assistant:
                compressed_data.append({
                    'original_user': user_text,
                    'compressed_user': compressed_user,
                    'original_assistant': assistant_text,
                    'compressed_assistant': compressed_assistant,
                    'compression_ratio_user': len(compressed_user) / len(user_text),
                    'compression_ratio_assistant': len(compressed_assistant) / len(assistant_text)
                })
                successful_compressions += 1
        
        # Save progress every 1000 examples
        if (i + 1) % 1000 == 0:
            print(f"💾 Processed {i+1} examples, {successful_compressions} successful compressions")
    
    # Save final compressed dataset
    print(f"💾 Saving {len(compressed_data)} compressed conversations to {output_file}")
    with open(output_file, 'w') as f:
        json.dump(compressed_data, f, indent=2)
    
    # Print statistics
    if compressed_data:
        avg_compression_user = sum(d['compression_ratio_user'] for d in compressed_data) / len(compressed_data)
        avg_compression_assistant = sum(d['compression_ratio_assistant'] for d in compressed_data) / len(compressed_data)
        
        print(f"\n📊 COMPRESSION STATISTICS")
        print(f"Total conversations: {len(compressed_data)}")
        print(f"Average user compression: {avg_compression_user:.2f}x")
        print(f"Average assistant compression: {avg_compression_assistant:.2f}x")
        processed_examples = i + 1
        print(f"Success rate: {successful_compressions}/{processed_examples} ({successful_compressions/processed_examples*100:.1f}%)")

def main():
    parser = argparse.ArgumentParser(description='Compress HuggingFace datasets to UTF-8 format')
    parser.add_argument('dataset', help='HuggingFace dataset name (e.g., HuggingFaceH4/ultrachat_200k)')
    parser.add_argument('--format', choices=['ultrachat', 'alpaca', 'sharegpt'], 
                       default='ultrachat', help='Dataset format')
    parser.add_argument('--output', '-o', default='compressed_dataset.json', 
                       help='Output file for compressed data')
    parser.add_argument('--max-examples', '-m', type=int, 
                       help='Maximum number of examples to process')
    parser.add_argument('--split', default='train', 
                       help='Dataset split to use (train, test, validation)')
    
    args = parser.parse_args()
    
    print(f"🚀 DATASET COMPRESSOR")
    print(f"Dataset: {args.dataset}")
    print(f"Format: {args.format}")
    print(f"Output: {args.output}")
    print(f"Split: {args.split}")
    if args.max_examples:
        print(f"Max examples: {args.max_examples}")
    print()
    
    process_dataset(
        dataset_name=args.dataset,
        dataset_format=args.format,
        output_file=args.output,
        max_examples=args.max_examples,
        split=args.split
    )

if __name__ == "__main__":
    main()