chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,117 @@
+"""
+This file specifies how MLC's BERT parameter maps from other formats, for example HuggingFace
+PyTorch, HuggingFace safetensors.
+"""
+
+import functools
+from typing import Literal
+
+import numpy as np
+
+from mlc_llm.loader import ExternMapping
+from mlc_llm.quantization import Quantization
+
+from .bert_model import BertConfig, BertModel
+
+
+def huggingface(
+    model_config: BertConfig,
+    quantization: Quantization,
+    hf_prefix: Literal["", "bert."] = "",
+) -> ExternMapping:
+    """Returns a parameter mapping that maps from the names of MLC LLM parameters to
+    the names of HuggingFace PyTorch parameters.
+
+    Parameters
+    ----------
+    model_config : BertConfig
+        The configuration of the BERT model.
+
+    quantization : Quantization
+        The quantization configuration.
+
+    hf_prefix : Literal["", "bert."]
+        Prefix used in HuggingFace weight names. Defaults to "" for standard
+        BERT models. Use "bert." for BGE models whose weights are prefixed.
+
+    Returns
+    -------
+    param_map : ExternMapping
+        The parameter mapping from MLC to HuggingFace PyTorch.
+    """
+    model = BertModel(model_config)
+    if quantization is not None:
+        model.to(quantization.model_dtype)
+    _, _named_params, _ = model.export_tvm(
+        spec=model.get_default_spec(),
+        allow_extern=True,
+    )
+    named_parameters = dict(_named_params)
+
+    mapping = ExternMapping()
+
+    def to_hf(name: str) -> str:
+        return f"{hf_prefix}{name}" if hf_prefix else name
+
+    for i in range(model_config.num_hidden_layers):
+        attn = f"encoder.layer.{i}.attention.self"
+        mlc_name = f"{attn}.qkv.weight"
+        mlc_param = named_parameters[mlc_name]
+        mapping.add_mapping(
+            mlc_name,
+            [
+                to_hf(f"{attn}.query.weight"),
+                to_hf(f"{attn}.key.weight"),
+                to_hf(f"{attn}.value.weight"),
+            ],
+            functools.partial(
+                lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
+                dtype=mlc_param.dtype,
+            ),
+        )
+
+        mlc_name = f"{attn}.qkv.bias"
+        mlc_param = named_parameters[mlc_name]
+        mapping.add_mapping(
+            mlc_name,
+            [
+                to_hf(f"{attn}.query.bias"),
+                to_hf(f"{attn}.key.bias"),
+                to_hf(f"{attn}.value.bias"),
+            ],
+            functools.partial(
+                lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
+                dtype=mlc_param.dtype,
+            ),
+        )
+
+    for mlc_name, mlc_param in named_parameters.items():
+        if mlc_name not in mapping.param_map:
+            mapping.add_mapping(
+                mlc_name,
+                [to_hf(mlc_name)],
+                functools.partial(
+                    lambda x, dtype: x.astype(dtype),
+                    dtype=mlc_param.dtype,
+                ),
+            )
+
+    # Mark unused weights that exist in HF but not in MLC
+    if hf_prefix:
+        mapping.add_unused(f"{hf_prefix}pooler.dense.weight")
+        mapping.add_unused(f"{hf_prefix}pooler.dense.bias")
+
+    return mapping
+
+
+def huggingface_bge(model_config: BertConfig, quantization: Quantization) -> ExternMapping:
+    """Returns a parameter mapping for BGE models.
+
+    BGE weights have no prefix but include extra unused weights:
+    pooler.dense.weight, pooler.dense.bias, embeddings.position_ids
+    """
+    mapping = huggingface(model_config, quantization, "")
+    mapping.add_unused("pooler.dense.weight")
+    mapping.add_unused("pooler.dense.bias")
+    mapping.add_unused("embeddings.position_ids")
+    return mapping