wehub-resource-sync
40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
from pathlib import Path
|
|
from typing import Union
|
|
|
|
import pytest
|
|
import tvm
|
|
|
|
from mlc_llm.loader import HuggingFaceLoader
|
|
from mlc_llm.model import MODEL_PRESETS, MODELS
|
|
from mlc_llm.quantization import QUANTIZATION
|
|
from mlc_llm.support import logging, tqdm
|
|
|
|
logging.enable_logging()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"param_path",
|
|
[
|
|
"./dist/models/llama-2-7b-w4-g128-awq.pt",
|
|
"./dist/models/Llama-2-7B-AWQ/model.safetensors",
|
|
],
|
|
)
|
|
def test_load_llama(param_path: Union[str, Path]):
|
|
path_params = Path(param_path)
|
|
|
|
model = MODELS["llama"]
|
|
quantization = QUANTIZATION["q4f16_awq"]
|
|
config = model.config.from_dict(MODEL_PRESETS["llama2_7b"])
|
|
loader = HuggingFaceLoader(
|
|
path=path_params,
|
|
extern_param_map=model.source["awq"](config, quantization),
|
|
)
|
|
with tqdm.redirect():
|
|
for _name, _param in loader.load(tvm.device("cpu")):
|
|
...
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_load_llama(param_path="./dist/models/llama-2-7b-w4-g128-awq.pt")
|
|
test_load_llama(param_path="./dist/models/Llama-2-7B-AWQ/model.safetensors")
|