Repository: mlc-ai/mlc-llm
File: python/mlc_llm/model/qwen3/qwen3_loader.py (Python)

Code before:
import functools
from typing import Callable, List

import numpy as np

from mlc_llm.loader import ExternMapping, QuantizeMapping
from mlc_llm.quantization import BlockScaleQuantize, Quantization

from .qwen3_model import Qwen3Config, Qwen3LMHeadModel


def huggingface(model_config: Qwen3Config, quantization: Quantization, hf_prefix: str = "model.") -> ExternMapping:
    """Returns a parameter mapping that maps from the names of MLC LLM parameters to
    the names of HuggingFace PyTorch parameters.

    Parameters
    ----------
    model_config : Qwen3Config
        The configuration of the Qwen3 model.

    quantization : Quantization
        The quantization configuration.

    hf_prefix : str
        Prefix used in HuggingFace weight names. Defaults to "model." for standard
        Qwen3 models. Use "" for Qwen3-Embedding models without prefix.

    Returns
    -------
    param_map : ExternMapping
        The parameter mapping from MLC to HuggingFace PyTorch.
    """
    model = Qwen3LMHeadModel(model_config)
    if quantization is not None:
        model.to(quantization.model_dtype)
    if isinstance(quantization, BlockScaleQuantize):
        model = quantization.quantize_model(model, QuantizeMapping({}, {}), "")
        if model_config.weight_block_size is None:
            raise ValueError(
                "The input Qwen3 model is not fp8 block quantized. "
                "Thus BlockScaleQuantize is not supported."
            )

    _, _named_params, _ = model.export_tvm(  # type: ignore[misc]
        spec=model.get_default_spec(),
        allow_extern=True,
    )
    named_parameters = dict(_named_params)

    mapping = ExternMapping()


Review comment (suggestion):
Actually, if there are only two options rather than being arbitrary strings, we can use type `Literal`.  For example, `hf_prefix: Literal["", "model."]`.

Code after:
import functools
from typing import Callable, List, Literal

import numpy as np

from mlc_llm.loader import ExternMapping, QuantizeMapping
from mlc_llm.quantization import BlockScaleQuantize, Quantization

from .qwen3_model import Qwen3Config, Qwen3LMHeadModel


def huggingface(
    model_config: Qwen3Config,
    quantization: Quantization,
    hf_prefix: Literal["", "model."] = "model.",
) -> ExternMapping:
    """Returns a parameter mapping that maps from the names of MLC LLM parameters to
    the names of HuggingFace PyTorch parameters.

    Parameters
    ----------
    model_config : Qwen3Config
        The configuration of the Qwen3 model.

    quantization : Quantization
        The quantization configuration.

    hf_prefix : Literal["", "model."]
        Prefix used in HuggingFace weight names. Defaults to "model." for standard
        Qwen3 models. Use "" for Qwen3-Embedding models without prefix.

    Returns
    -------
    param_map : ExternMapping
        The parameter mapping from MLC to HuggingFace PyTorch.
    """
    model = Qwen3LMHeadModel(model_config)
    if quantization is not None:
        model.to(quantization.model_dtype)
    if isinstance(quantization, BlockScaleQuantize):
        model = quantization.quantize_model(model, QuantizeMapping({}, {}), "")
        if model_config.weight_block_size is None:
            raise ValueError(
                "The input Qwen3 model is not fp8 block quantized. "
                "Thus BlockScaleQuantize is not supported."
            )

    _, _named_params, _ = model.export_tvm(  # type: ignore[misc]
        spec=model.get_default_spec(),
        allow_extern=True,
    )