setup_mixassist.py

#!/usr/bin/env python3
"""
MixAssist Dataset Setup Script
Downloads and configures the MixAssist dataset for use with Carla MCP Server
"""
import argparse
import os
import sys
from pathlib import Path
from typing import Optional


def check_requirements():
    """Check if required packages are installed"""
    try:
        import datasets
        import pandas as pd
    except ImportError as e:
        print(f"❌ Missing required package: {e.name}")
        print("\nInstall required packages:")
        print("  pip install datasets pandas pyarrow")
        sys.exit(1)


def download_dataset(output_dir: Path, force: bool = False) -> bool:
    """Download MixAssist dataset from Hugging Face

    Args:
        output_dir: Directory to store the dataset
        force: Force redownload even if dataset exists

    Returns:
        True if successful, False otherwise
    """
    from datasets import load_dataset

    # Check if dataset already exists
    if output_dir.exists() and not force:
        parquet_files = list(output_dir.glob("*.parquet"))
        if len(parquet_files) >= 3:
            print(f"✅ Dataset already exists at {output_dir}")
            print(f"   Found {len(parquet_files)} parquet files")
            print("\n   Use --force to redownload")
            return True

    # Create output directory
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"📥 Downloading MixAssist dataset from Hugging Face...")
    print(f"   Destination: {output_dir}")

    try:
        # Load dataset from Hugging Face
        # Dataset: https://huggingface.co/datasets/MixAssist/mixassist
        dataset = load_dataset("MixAssist/mixassist", trust_remote_code=True)

        print(f"\n✅ Downloaded dataset with {len(dataset)} splits:")
        for split_name, split_data in dataset.items():
            print(f"   - {split_name}: {len(split_data)} conversations")

            # Save each split as parquet
            output_file = output_dir / f"{split_name}-00000-of-00001.parquet"
            split_data.to_parquet(str(output_file))
            print(f"   💾 Saved: {output_file.name}")

        print(f"\n🎉 Dataset successfully downloaded to: {output_dir}")
        return True

    except Exception as e:
        print(f"\n❌ Failed to download dataset: {e}")
        print("\nTroubleshooting:")
        print("  1. Check your internet connection")
        print("  2. Verify Hugging Face access (might need login for some datasets)")
        print("  3. Try: huggingface-cli login")
        return False


def create_config(dataset_path: Path, config_file: Optional[Path] = None) -> bool:
    """Create configuration file with dataset path

    Args:
        dataset_path: Path to the dataset directory
        config_file: Path to config file (default: .env in project root)

    Returns:
        True if successful
    """
    if config_file is None:
        config_file = Path(__file__).parent / ".env"

    # Check if config already exists
    if config_file.exists():
        print(f"\n⚠️  Config file already exists: {config_file}")
        response = input("   Overwrite? (y/N): ").strip().lower()
        if response != 'y':
            print("   Skipped config creation")
            return True

    # Create config content
    config_content = f"""# Carla MCP Server Configuration
# Auto-generated by setup_mixassist.py

# MixAssist Dataset Configuration
MIXASSIST_DATASET_PATH={dataset_path.absolute()}

# Optional: Enable/disable MixAssist resources
MIXASSIST_ENABLED=true
"""

    try:
        config_file.write_text(config_content)
        print(f"\n✅ Created config file: {config_file}")
        print(f"   Dataset path: {dataset_path.absolute()}")
        return True
    except Exception as e:
        print(f"\n❌ Failed to create config: {e}")
        return False


def verify_dataset(dataset_path: Path) -> bool:
    """Verify that the dataset is valid and complete

    Args:
        dataset_path: Path to verify

    Returns:
        True if valid, False otherwise
    """
    import pandas as pd

    print(f"\n🔍 Verifying dataset at: {dataset_path}")

    if not dataset_path.exists():
        print(f"   ❌ Directory does not exist")
        return False

    # Check for required parquet files
    required_splits = ["train", "test", "validation"]
    total_conversations = 0

    for split in required_splits:
        parquet_file = dataset_path / f"{split}-00000-of-00001.parquet"
        if not parquet_file.exists():
            print(f"   ❌ Missing {split} split: {parquet_file.name}")
            return False

        try:
            df = pd.read_parquet(parquet_file)
            conversation_count = len(df)
            total_conversations += conversation_count
            print(f"   ✅ {split}: {conversation_count} conversations")
        except Exception as e:
            print(f"   ❌ Failed to read {split} split: {e}")
            return False

    print(f"\n✅ Dataset verification passed!")
    print(f"   Total: {total_conversations} conversations across {len(required_splits)} splits")

    return True


def main():
    parser = argparse.ArgumentParser(
        description="Setup MixAssist dataset for Carla MCP Server",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Download to default location and create config
  python setup_mixassist.py --download

  # Download to custom location
  python setup_mixassist.py --download --output ~/datasets/mixassist

  # Just create config for existing dataset
  python setup_mixassist.py --path /path/to/existing/dataset

  # Verify existing dataset
  python setup_mixassist.py --verify --path /path/to/dataset

  # Force redownload
  python setup_mixassist.py --download --force
"""
    )

    parser.add_argument(
        "--download",
        action="store_true",
        help="Download the dataset from Hugging Face"
    )

    parser.add_argument(
        "--output",
        type=Path,
        default=Path.home() / ".cache" / "mixassist" / "data",
        help="Output directory for downloaded dataset (default: ~/.cache/mixassist/data)"
    )

    parser.add_argument(
        "--path",
        type=Path,
        help="Path to existing dataset (for config creation or verification)"
    )

    parser.add_argument(
        "--verify",
        action="store_true",
        help="Verify dataset integrity"
    )

    parser.add_argument(
        "--force",
        action="store_true",
        help="Force redownload even if dataset exists"
    )

    parser.add_argument(
        "--config",
        type=Path,
        help="Path to config file (default: .env in project root)"
    )

    parser.add_argument(
        "--no-config",
        action="store_true",
        help="Skip config file creation"
    )

    args = parser.parse_args()

    # Show header
    print("=" * 60)
    print("MixAssist Dataset Setup for Carla MCP Server")
    print("=" * 60)

    # Determine dataset path
    dataset_path = args.path if args.path else args.output

    # Download if requested
    if args.download:
        check_requirements()
        if not download_dataset(args.output, force=args.force):
            sys.exit(1)
        dataset_path = args.output

    # Verify if requested or after download
    if args.verify or args.download:
        check_requirements()
        if not verify_dataset(dataset_path):
            print("\n❌ Dataset verification failed")
            sys.exit(1)

    # Create config unless explicitly disabled
    if not args.no_config:
        if not dataset_path.exists():
            print(f"\n❌ Dataset path does not exist: {dataset_path}")
            print("   Run with --download to download the dataset first")
            sys.exit(1)

        if not create_config(dataset_path, args.config):
            sys.exit(1)

    # Show success message
    print("\n" + "=" * 60)
    print("🎉 Setup Complete!")
    print("=" * 60)
    print("\nNext steps:")
    print("  1. The MixAssist resources will be automatically available in the MCP server")
    print("  2. Restart your MCP server if it's already running")
    print("  3. Resources are accessible via URIs like: mixassist://index")
    print("\nFor more information, see: CLAUDE.md")
    print("=" * 60)


if __name__ == "__main__":
    main()