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,150 @@
+import json
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from mlc_llm.support import logging
+from mlc_llm.support.auto_weight import detect_weight
+
+logging.enable_logging()
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+def _create_json_file(json_path, data):
+    with open(json_path, "w", encoding="utf-8") as i_f:
+        json.dump(data, i_f)
+
+
+@pytest.mark.parametrize(
+    "weight_format, index_filename, result",
+    [
+        ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"),
+        (
+            "huggingface-safetensor",
+            "model.safetensors.index.json",
+            "huggingface-safetensor",
+        ),
+        ("auto", "pytorch_model.bin.index.json", "huggingface-torch"),
+        ("auto", "model.safetensors.index.json", "huggingface-safetensor"),
+    ],
+)
+def test_detect_weight(weight_format, index_filename, result):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        if index_filename is not None:
+            weight_index_file = base_path / index_filename
+            _create_json_file(weight_index_file, {})
+        assert detect_weight(base_path, None, weight_format) == (
+            weight_index_file,
+            result,
+        )
+
+
+@pytest.mark.parametrize(
+    "weight_format, index_filename, result",
+    [
+        ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"),
+        (
+            "huggingface-safetensor",
+            "model.safetensors.index.json",
+            "huggingface-safetensor",
+        ),
+        ("auto", "pytorch_model.bin.index.json", "huggingface-torch"),
+        ("auto", "model.safetensors.index.json", "huggingface-safetensor"),
+    ],
+)
+def test_detect_weight_in_config_json(weight_format, index_filename, result):
+    with (
+        tempfile.TemporaryDirectory() as config_dir,
+        tempfile.TemporaryDirectory() as weight_dir,
+    ):
+        config_path = Path(config_dir)
+        weight_path = Path(weight_dir)
+        config_json_path = config_path / "config.json"
+        _create_json_file(config_json_path, {"weight_path": weight_dir})
+        if index_filename is not None:
+            weight_index_file = weight_path / index_filename
+            _create_json_file(weight_index_file, {})
+
+        assert detect_weight(None, config_json_path, weight_format) == (
+            weight_index_file,
+            result,
+        )
+
+
+@pytest.mark.parametrize(
+    "weight_format, index_filename, result",
+    [
+        ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"),
+        (
+            "huggingface-safetensor",
+            "model.safetensors.index.json",
+            "huggingface-safetensor",
+        ),
+        ("auto", "pytorch_model.bin.index.json", "huggingface-torch"),
+        ("auto", "model.safetensors.index.json", "huggingface-safetensor"),
+    ],
+)
+def test_detect_weight_same_dir_config_json(weight_format, index_filename, result):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        config_json_path = base_path / "config.json"
+        _create_json_file(config_json_path, {})
+        if index_filename is not None:
+            weight_index_file = Path(os.path.join(tmpdir, index_filename))
+            _create_json_file(weight_index_file, {})
+        assert detect_weight(None, config_json_path, weight_format) == (
+            weight_index_file,
+            result,
+        )
+
+
+def test_find_weight_fail():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        with pytest.raises(ValueError):
+            detect_weight(Path("do/not/exist"), base_path, "awq")
+    with pytest.raises(AssertionError):
+        detect_weight(None, Path("do/not/exist"), "awq")
+
+
+if __name__ == "__main__":
+    test_detect_weight("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch")
+    test_detect_weight(
+        "huggingface-safetensor",
+        "model.safetensors.index.json",
+        "huggingface-safetensor",
+    )
+    test_detect_weight("auto", "pytorch_model.bin.index.json", "huggingface-torch")
+    test_detect_weight("auto", "model.safetensors.index.json", "huggingface-safetensor")
+    test_detect_weight_in_config_json(
+        "huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"
+    )
+    test_detect_weight_in_config_json(
+        "huggingface-safetensor",
+        "model.safetensors.index.json",
+        "huggingface-safetensor",
+    )
+    test_detect_weight_in_config_json("auto", "pytorch_model.bin.index.json", "huggingface-torch")
+    test_detect_weight_in_config_json(
+        "auto", "model.safetensors.index.json", "huggingface-safetensor"
+    )
+    test_detect_weight_same_dir_config_json(
+        "huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"
+    )
+    test_detect_weight_same_dir_config_json(
+        "huggingface-safetensor",
+        "model.safetensors.index.json",
+        "huggingface-safetensor",
+    )
+    test_detect_weight_same_dir_config_json(
+        "auto", "pytorch_model.bin.index.json", "huggingface-torch"
+    )
+    test_detect_weight_same_dir_config_json(
+        "auto", "model.safetensors.index.json", "huggingface-safetensor"
+    )
+    test_find_weight_fail()