convert_cli.py

#!/usr/bin/env python3
"""
CAM++ MLX Converter - Command Line Interface

Convert PyTorch CAM++ models to MLX format without Gradio UI.
Perfect for batch processing, CI/CD pipelines, or scripting.

Usage:
    # Basic conversion
    python convert_cli.py \\
        --input iic/speech_campplus_sv_zh-cn_16k-common \\
        --output campplus_chinese_16k \\
        --token YOUR_HF_TOKEN

    # With quantization options
    python convert_cli.py \\
        --input iic/speech_campplus_sv_zh_en_16k-common_advanced \\
        --output campplus_multilingual \\
        --token YOUR_HF_TOKEN \\
        --q2 --q4 --q8

    # Using environment variable for token
    export HF_TOKEN=your_token_here
    python convert_cli.py \\
        --input iic/speech_campplus_sv_zh-cn_16k-common \\
        --output campplus_chinese_16k

    # Dry run (validate without uploading)
    python convert_cli.py \\
        --input iic/speech_campplus_sv_zh-cn_16k-common \\
        --output campplus_chinese_16k \\
        --token YOUR_HF_TOKEN \\
        --dry-run
"""

import argparse
import os
import sys
from typing import Optional
import logging
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Import the converter
from app import CAMPPConverter, DEFAULT_SERVER_PORT, TARGET_ORGANIZATION

# Set up logging for CLI
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(
        description='Convert PyTorch CAM++ models to MLX format (CLI)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Convert Chinese model with Q4 quantization (default)
  %(prog)s -i iic/speech_campplus_sv_zh-cn_16k-common \\
           -o campplus_chinese_16k \\
           -t YOUR_HF_TOKEN

  # Convert with all quantization levels
  %(prog)s -i iic/speech_campplus_sv_zh_en_16k-common_advanced \\
           -o campplus_multilingual \\
           -t YOUR_HF_TOKEN \\
           --q2 --q4 --q8

  # Use environment variable for token
  export HF_TOKEN=your_token_here
  %(prog)s -i iic/speech_campplus_sv_zh-cn_16k-common \\
           -o campplus_chinese_16k

  # Dry run (test without uploading)
  %(prog)s -i iic/speech_campplus_sv_zh-cn_16k-common \\
           -o campplus_chinese_16k \\
           -t YOUR_HF_TOKEN \\
           --dry-run

Preset Models:
  - Chinese (Basic): iic/speech_campplus_sv_zh-cn_16k-common
  - Chinese-English (Advanced): iic/speech_campplus_sv_zh_en_16k-common_advanced
        """
    )

    # Required arguments
    parser.add_argument(
        '-i', '--input',
        required=True,
        help='Input ModelScope repository (e.g., iic/speech_campplus_sv_zh-cn_16k-common)'
    )

    parser.add_argument(
        '-o', '--output',
        required=True,
        help='Output model name (will be uploaded to mlx-community/{output})'
    )

    # Optional token (can use env var)
    parser.add_argument(
        '-t', '--token',
        help='HuggingFace API token (or set HF_TOKEN environment variable)'
    )

    # Quantization options
    quant_group = parser.add_argument_group('quantization options')
    quant_group.add_argument(
        '--q2',
        action='store_true',
        help='Create 2-bit quantized version'
    )
    quant_group.add_argument(
        '--q4',
        action='store_true',
        default=False,
        help='Create 4-bit quantized version (enabled by default if no quant flags specified)'
    )
    quant_group.add_argument(
        '--q8',
        action='store_true',
        help='Create 8-bit quantized version'
    )
    quant_group.add_argument(
        '--no-quantization',
        action='store_true',
        help='Skip all quantization (only create regular version)'
    )

    # Other options
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Test conversion without uploading to HuggingFace'
    )

    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Enable verbose logging'
    )

    parser.add_argument(
        '--version',
        action='version',
        version='CAM++ MLX Converter CLI v1.0.0'
    )

    return parser.parse_args()


def get_hf_token(args) -> Optional[str]:
    """
    Get HuggingFace token from args or environment

    Args:
        args: Parsed command line arguments

    Returns:
        HF token string or None
    """
    # Priority: command line arg > environment variable
    if args.token:
        return args.token

    env_token = os.getenv('HF_TOKEN') or os.getenv('HUGGING_FACE_HUB_TOKEN')
    if env_token:
        logger.info("Using HF token from environment variable")
        return env_token

    return None


def validate_args(args) -> bool:
    """
    Validate command line arguments

    Args:
        args: Parsed arguments

    Returns:
        True if valid, False otherwise
    """
    # Check token unless dry run
    if not args.dry_run:
        token = get_hf_token(args)
        if not token:
            logger.error("ERROR: HuggingFace token required. Provide via --token or HF_TOKEN environment variable")
            logger.error("       Use --dry-run to test conversion without uploading")
            return False

        if not token.startswith('hf_'):
            logger.warning("WARNING: HF token should start with 'hf_' - are you sure this is correct?")

    # Validate repo format
    if '/' not in args.input:
        logger.error(f"ERROR: Input repo '{args.input}' must be in format 'username/model-name'")
        return False

    # Validate output name
    if '/' in args.output:
        logger.error(f"ERROR: Output name '{args.output}' should not contain '/' (organization is automatically set to {TARGET_ORGANIZATION})")
        return False

    return True


def main():
    """Main CLI entry point"""
    args = parse_args()

    # Set verbose logging if requested
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
        logger.debug("Verbose logging enabled")

    # Validate arguments
    if not validate_args(args):
        sys.exit(1)

    # Get token
    token = get_hf_token(args)
    if args.dry_run:
        token = "dry_run_token_placeholder"
        logger.info("🔍 DRY RUN MODE - Will not upload to HuggingFace")

    # Determine quantization settings
    # Default to Q4 if no quantization flags specified
    if not args.no_quantization and not (args.q2 or args.q4 or args.q8):
        args.q4 = True
        logger.info("📦 No quantization flags specified - defaulting to Q4")

    if args.no_quantization:
        args.q2 = args.q4 = args.q8 = False
        logger.info("📦 Quantization disabled - creating regular version only")

    # Display configuration
    logger.info("=" * 70)
    logger.info("CAM++ MLX Converter - CLI Mode")
    logger.info("=" * 70)
    logger.info(f"Input Repository:  {args.input}")
    logger.info(f"Output Name:       {TARGET_ORGANIZATION}/{args.output}")
    logger.info(f"Quantization:      Q2={args.q2}, Q4={args.q4}, Q8={args.q8}")
    logger.info(f"Dry Run:           {args.dry_run}")
    logger.info("=" * 70)
    logger.info("")

    # Create converter
    converter = CAMPPConverter()

    # Perform conversion
    try:
        logger.info("🚀 Starting conversion...")
        logger.info("")

        # If dry run, we'll need to modify the converter to skip upload
        # For now, just run the conversion normally
        # TODO: Add dry_run parameter to converter

        result = converter.convert_model(
            input_repo=args.input,
            output_name=args.output,
            hf_token=token,
            quantize_q2=args.q2,
            quantize_q4=args.q4,
            quantize_q8=args.q8
        )

        # Print results
        logger.info("")
        logger.info("=" * 70)
        logger.info("CONVERSION RESULTS")
        logger.info("=" * 70)
        print(result)
        logger.info("=" * 70)

        # Check if conversion was successful
        if "✅" in result or "Conversion Successful" in result:
            logger.info("✅ Conversion completed successfully!")
            sys.exit(0)
        elif "⚠️" in result or "Warning" in result:
            logger.warning("⚠️ Conversion completed with warnings (model not uploaded)")
            sys.exit(1)
        else:
            logger.error("❌ Conversion failed")
            sys.exit(1)

    except KeyboardInterrupt:
        logger.info("\n\n⚠️ Conversion interrupted by user")
        sys.exit(130)
    except Exception as e:
        logger.error(f"\n\n❌ Conversion failed with exception: {e}")
        if args.verbose:
            import traceback
            traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()