diff --git a/conversion/__init__.py b/conversion/__init__.py index 02102fac8d..465c010f66 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -31,6 +31,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "BertForSequenceClassification": "bert", "BertModel": "bert", "BitnetForCausalLM": "bitnet", + "BitNetForCausalLM": "bitnet", "BloomForCausalLM": "bloom", "BloomModel": "bloom", "CamembertModel": "bert", diff --git a/conversion/bitnet.py b/conversion/bitnet.py index a66446abee..0c2baee876 100644 --- a/conversion/bitnet.py +++ b/conversion/bitnet.py @@ -8,7 +8,7 @@ if TYPE_CHECKING: from .base import ModelBase, TextModel, gguf -@ModelBase.register("BitnetForCausalLM") +@ModelBase.register("BitnetForCausalLM", "BitNetForCausalLM") class BitnetModel(TextModel): model_arch = gguf.MODEL_ARCH.BITNET