dflash_mlx/convert.py

"""
Convert PyTorch DFlash drafter models to MLX format.

Handles weight conversion from PyTorch safetensors to MLX arrays,
compatible with any z-lab DFlash drafter.

Updated to work with the universal adapter system for any target model family.
"""

import json
import os
from pathlib import Path
from typing import Optional, Dict, Tuple
import mlx.core as mx
from transformers import AutoConfig
from huggingface_hub import hf_hub_download, snapshot_download


def _convert_key(key: str) -> str:
    """Convert PyTorch parameter names to MLX format.
    
    Handles various naming conventions across model families.
    """
    # Replace PyTorch-specific prefixes
    key = key.replace("model.", "")
    # Standardize naming
    replacements = {
        "embed_tokens": "embed_tokens",
        "layers.": "layers.",
        "self_attn.": "self_attn.",
        "mlp.": "mlp.",
        "input_layernorm": "input_layernorm",
        "post_attention_layernorm": "post_attention_layernorm",
        "norm": "norm",
        "lm_head": "lm_head",
        "q_proj": "q_proj",
        "k_proj": "k_proj",
        "v_proj": "v_proj",
        "o_proj": "o_proj",
        "gate_proj": "gate_proj",
        "up_proj": "up_proj",
        "down_proj": "down_proj",
        "fc": "fc",
        "hidden_norm": "hidden_norm",
        "q_norm": "q_norm",
        "k_norm": "k_norm",
        "weight": "weight",
    }
    return key


def _transpose_if_needed(key: str, tensor) -> mx.array:
    """Transpose linear layer weights from PyTorch to MLX format.
    
    Linear layers in PyTorch are [out, in], MLX expects [in, out].
    """
    if "proj" in key or "fc" in key or "lm_head" in key or "embed" in key:
        if len(tensor.shape) == 2:
            return mx.array(tensor.T)
    return mx.array(tensor)


def convert_dflash_to_mlx(
    pytorch_model_id: str,
    output_path: str,
    trust_remote_code: bool = True,
    token: Optional[str] = None,
) -> str:
    """Convert a PyTorch DFlash drafter to MLX format.
    
    Args:
        pytorch_model_id: Hugging Face model ID (e.g., "z-lab/Qwen3-4B-DFlash-b16")
        output_path: Local directory to save converted model
        trust_remote_code: Whether to trust custom modeling code
        token: HF API token for gated/private models
    
    Returns:
        Path to the converted model directory
    """
    output_path = Path(output_path)
    output_path.mkdir(parents=True, exist_ok=True)

    print(f"[Convert] Downloading {pytorch_model_id}...")
    
    # Download model files
    repo_path = snapshot_download(
        repo_id=pytorch_model_id,
        token=token,
        ignore_patterns=["*.md", "*.png", "*.jpg", "*.gif", "*.jpeg"],
    )
    repo_path = Path(repo_path)

    # Load PyTorch model config
    print("[Convert] Loading PyTorch config...")
    config = AutoConfig.from_pretrained(
        repo_path,
        trust_remote_code=trust_remote_code,
    )

    # Extract DFlash-specific config
    dflash_config = {
        "vocab_size": getattr(config, "vocab_size", 151936),
        "hidden_size": getattr(config, "hidden_size", 1024),
        "num_hidden_layers": getattr(config, "num_hidden_layers", 5),
        "num_attention_heads": getattr(config, "num_attention_heads", 16),
        "num_key_value_heads": getattr(config, "num_key_value_heads", 4),
        "intermediate_size": getattr(config, "intermediate_size", 2816),
        "max_position_embeddings": getattr(config, "max_position_embeddings", 32768),
        "rms_norm_eps": getattr(config, "rms_norm_eps", 1e-6),
        "block_size": getattr(config, "block_size", 16),
        "rope_base": getattr(config, "rope_theta", 10000.0),
    }
    
    # Extract target layer IDs if present in config
    if hasattr(config, "target_layer_ids"):
        dflash_config["target_layer_ids"] = config.target_layer_ids
    elif hasattr(config, "dflash_config") and hasattr(config.dflash_config, "target_layer_ids"):
        dflash_config["target_layer_ids"] = config.dflash_config.target_layer_ids

    # Load weights from safetensors
    print("[Convert] Loading weights from safetensors...")
    try:
        from safetensors.torch import load_file
        
        # Find all safetensors files
        safetensors_files = sorted(repo_path.glob("*.safetensors"))
        
        if safetensors_files:
            pt_weights = {}
            for st_file in safetensors_files:
                print(f"  Loading {st_file.name}...")
                partial = load_file(str(st_file))
                pt_weights.update(partial)
        else:
            # Try pytorch_model.bin
            bin_file = repo_path / "pytorch_model.bin"
            if bin_file.exists():
                import torch
                pt_weights = torch.load(str(bin_file), map_location="cpu")
            else:
                raise FileNotFoundError("No safetensors or pytorch_model.bin found")
    except ImportError:
        # Fallback to torch load
        import torch
        weights_file = repo_path / "pytorch_model.bin"
        if weights_file.exists():
            pt_weights = torch.load(str(weights_file), map_location="cpu")
        else:
            raise FileNotFoundError("No weight files found and safetensors not installed")

    # Convert weights
    print(f"[Convert] Converting {len(pt_weights)} parameters...")
    mlx_weights = {}
    for key, tensor in pt_weights.items():
        mlx_key = _convert_key(key)
        mlx_weights[mlx_key] = _transpose_if_needed(key, tensor)

    # Save MLX weights (try safetensors, fallback to npz)
    weights_path = output_path / "weights.npz"
    try:
        # Use numpy format if safetensors save is problematic
        import numpy as np
        np_weights = {k: np.array(v) for k, v in mlx_weights.items()}
        np.savez(str(weights_path), **np_weights)
        print(f"[Convert] Saved weights to {weights_path}")
    except Exception as e:
        print(f"[Convert] Warning: Could not save weights: {e}")
        # Try direct mlx save
        try:
            mx.savez(str(weights_path), **mlx_weights)
        except Exception as e2:
            print(f"[Convert] Error saving weights: {e2}")
            raise

    # Save config
    config_path = output_path / "config.json"
    with open(config_path, "w") as f:
        json.dump(dflash_config, f, indent=2)

    # Save target model mapping
    target_info = {
        "source_model": pytorch_model_id,
        "target_model": infer_target_model(pytorch_model_id),
        "conversion_date": str(Path(__file__).stat().st_mtime),
    }
    info_path = output_path / "model_info.json"
    with open(info_path, "w") as f:
        json.dump(target_info, f, indent=2)

    print(f"[Convert] Done! Model saved to {output_path}")
    print(f"  Config: {dflash_config}")
    print(f"  Target: {target_info['target_model']}")
    return str(output_path)


def infer_target_model(dflash_model_id: str) -> str:
    """Infer the target model from DFlash drafter ID.
    
    Maps known drafter checkpoints to their corresponding target models.
    Supports all official z-lab DFlash models plus community variants.
    """
    # Map drafter IDs to target models
    mapping = {
        # Qwen3 series
        "Qwen3-4B-DFlash": "Qwen/Qwen3-4B",
        "Qwen3-8B-DFlash": "Qwen/Qwen3-8B",
        "Qwen3-32B-DFlash": "Qwen/Qwen3-32B",
        # Qwen3.5 series
        "Qwen3.5-4B-DFlash": "Qwen/Qwen3.5-4B",
        "Qwen3.5-9B-DFlash": "Qwen/Qwen3.5-9B",
        "Qwen3.5-27B-DFlash": "Qwen/Qwen3.5-27B",
        "Qwen3.5-35B-A3B-DFlash": "Qwen/Qwen3.5-35B-A3B",
        "Qwen3.5-122B-A10B-DFlash": "Qwen/Qwen3.5-122B-A10B",
        # Qwen3.6 series
        "Qwen3.6-27B-DFlash": "Qwen/Qwen3.6-27B",
        "Qwen3.6-35B-A3B-DFlash": "Qwen/Qwen3.6-35B-A3B",
        # Qwen Coder
        "Qwen3-Coder-Next-DFlash": "Qwen/Qwen3-Coder-Next",
        "Qwen3-Coder-30B-A3B-DFlash": "Qwen/Qwen3-Coder-30B-A3B",
        # LLaMA
        "LLaMA3.1-8B-Instruct-DFlash": "meta-llama/Llama-3.1-8B-Instruct",
        "LLaMA3.1-70B-Instruct-DFlash": "meta-llama/Llama-3.1-70B-Instruct",
        # Gemma
        "gemma-4-31B-it-DFlash": "google/gemma-4-31b-it",
        "gemma-4-26B-A4B-it-DFlash": "google/gemma-4-26b-a4b-it",
        # GPT-OSS
        "gpt-oss-20b-DFlash": "openai/gpt-oss-20b",
        "gpt-oss-120b-DFlash": "openai/gpt-oss-120b",
        # Kimi
        "Kimi-K2.5-DFlash": "moonshotai/Kimi-K2.5",
        # MiniMax
        "MiniMax-M2.5-DFlash": "MiniMax/MiniMax-M2.5",
    }
    
    # Direct mapping lookup
    for key, target in mapping.items():
        if key in dflash_model_id:
            return target
    
    # Generic inference by model family
    if "Qwen3.6" in dflash_model_id:
        return "Qwen/Qwen3.6-27B"
    elif "Qwen3.5" in dflash_model_id:
        return "Qwen/Qwen3.5-9B"
    elif "Qwen3-Coder" in dflash_model_id:
        return "Qwen/Qwen3-Coder-Next"
    elif "Qwen3" in dflash_model_id:
        return "Qwen/Qwen3-4B"
    elif "LLaMA" in dflash_model_id or "Llama" in dflash_model_id or "llama" in dflash_model_id:
        return "meta-llama/Llama-3.1-8B-Instruct"
    elif "gemma" in dflash_model_id.lower():
        return "google/gemma-4-31b-it"
    elif "gpt-oss" in dflash_model_id.lower():
        return "openai/gpt-oss-20b"
    elif "Kimi" in dflash_model_id:
        return "moonshotai/Kimi-K2.5"
    elif "MiniMax" in dflash_model_id:
        return "MiniMax/MiniMax-M2.5"
    
    return "unknown"


def load_mlx_dflash(
    model_path: str,
) -> Tuple:
    """Load a converted MLX DFlash model.
    
    Args:
        model_path: Path to converted MLX model directory
    
    Returns:
        Tuple of (model, config)
    """
    from .model import DFlashDraftModel

    model_path = Path(model_path)
    
    # Load config
    with open(model_path / "config.json", "r") as f:
        config = json.load(f)

    # Load weights
    weights_path = model_path / "weights.npz"
    if not weights_path.exists():
        # Try alternative extensions
        for ext in [".safetensors", ".mlx", ".npz"]:
            alt = model_path / f"weights{ext}"
            if alt.exists():
                weights_path = alt
                break
    
    if not weights_path.exists():
        raise FileNotFoundError(f"No weights found in {model_path}")
    
    weights = mx.load(str(weights_path))
    
    # Build model
    model = DFlashDraftModel(
        vocab_size=config["vocab_size"],
        hidden_size=config["hidden_size"],
        num_layers=config["num_hidden_layers"],
        num_heads=config["num_attention_heads"],
        num_kv_heads=config["num_key_value_heads"],
        intermediate_size=config["intermediate_size"],
        max_seq_len=config["max_position_embeddings"],
        block_size=config.get("block_size", 16),
        rope_base=config.get("rope_base", 10000.0),
        target_layer_ids=config.get("target_layer_ids", None),
    )

    # Load weights into model
    model.update(weights)

    return model, config


def main():
    """CLI entry point for conversion."""
    import argparse
    parser = argparse.ArgumentParser(description="Convert PyTorch DFlash drafter to MLX")
    parser.add_argument("--model", required=True, help="HF model ID of PyTorch drafter")
    parser.add_argument("--output", required=True, help="Output directory")
    parser.add_argument("--trust-remote-code", action="store_true", default=True)
    parser.add_argument("--token", default=None, help="HF token for gated models")
    args = parser.parse_args()
    
    convert_dflash_to_mlx(
        pytorch_model_id=args.model,
        output_path=args.output,
        trust_remote_code=args.trust_remote_code,
        token=args.token,
    )


if __name__ == "__main__":
    main()