ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix quant tag mapping

Browse source
This commit is contained in:
jeromeku 2025-03-30 16:14:33 -07:00
commit 7db0679f72
1 changed files with 4 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

14
unsloth/registry/registry.py
View file

@ -13,13 +13,12 @@ BNB_QUANTIZED_TAG = "bnb-4bit"
UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG
GGUF_TAG = "GGUF"
QUANT_TYPE_MAP = {
QUANT_TAG_MAP = {
QuantType.BNB: BNB_QUANTIZED_TAG,
QuantType.UNSLOTH: UNSLOTH_DYNAMIC_QUANT_TAG,
QuantType.GGUF: GGUF_TAG,
QuantType.NONE: None,
}
QUANT_TYPES = list(QUANT_TYPE_MAP.keys())
@dataclass
@ -52,13 +51,8 @@ class ModelInfo:
def append_quant_type(
key: str, quant_type: QuantType = None
):
if quant_type:
if quant_type == "bnb":
key = "-".join([key, QUANT_TYPE_MAP["bnb"]])
elif quant_type == "unsloth":
key = "-".join([key, QUANT_TYPE_MAP["unsloth"]])
elif quant_type == "GGUF":
key = "-".join([key, QUANT_TYPE_MAP["GGUF"]])
if quant_type != QuantType.NONE:
key = "-".join([key, QUANT_TAG_MAP[quant_type]])
return key
@classmethod
@ -108,7 +102,7 @@ def register_model(
key = f"{org}/{name}"
if key in MODEL_REGISTRY:
raise ValueError(f"Model {key} already registered")
raise ValueError(f"Model {key} already registered, current keys: {MODEL_REGISTRY.keys()}")
MODEL_REGISTRY[key] = model_info_cls(
org=org,