Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-08-26 06:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
157b81fe6dbfec7d7ce91ef7cd9c6bc0c218d6fe
koboldcpp/conversion/__init__.py
T
Bar Haim 157b81fe6d model : Granite-Switch Architecture (#25107) 
* granite-switch: add llama.cpp backend (POC, CPU)

New "granite-switch" architecture: a dense, all-attention Granite-4.1
model with N embedded LoRA adapters selected per-token by control tokens.

- gguf-py schema (arch, KV keys, stacked LoRA tensor names) + writer helpers
- conversion/granite.py: GraniteSwitchModel converter (stacks N adapters +
  zero base slot into per-projection A/B tensors; emits switch metadata)
- C++ arch registration (llama-arch.{h,cpp}, llama-model.{h,cpp})
- src/models/granite_switch.cpp: load + per-token switched-LoRA graph via
  ggml_mul_mat_id over stacked tensors; sticky per-token index + control-token
  substitution in llm_graph_input_switch::set_input
- llm_graph_input_switch in src/models/models.h

Runs end-to-end on CPU: convert 3b checkpoint (842 tensors, stacked dim 13)
and generate on both base and control-token paths. Sticky switch state is
single-sequence (POC); full multi-sequence machinery is a follow-up.

* granite-switch: add Mac (Metal) build + mid-sequence switch demo script

Self-contained script to build llama.cpp on Apple Silicon (Metal),
convert the composed 3b checkpoint, and run the crisp mid-sequence
adapter-switch demos verified on Vela:
  - answerability: <|answerability|> mid-seq -> "unanswerable"
  - query_rewrite: <|query_rewrite|> mid-seq -> {"rewritten_question": ...}
Each demo runs the same prompt twice, differing only by a control token
placed before the assistant turn, so the per-token switch is visible.

* granite-switch mac demo: add -no-cnv so each run is one-shot

The composed model ships a chat template, so llama-completion auto-enables
interactive conversation mode and halts at a `>` prompt after generating,
stalling the script. -no-cnv disables conversation mode: generate once from
the raw prompt and exit (also prints special tokens, making the switch visible).

* granite-switch: replace global sticky index with in-graph router attention

The POC computed the per-token adapter index on the CPU and carried it
across ubatches in ONE global `mutable int32_t poc_sticky_index`, reset
only when a ubatch contained sequence position 0. That global had two
problems:

  1. Concurrency: with multiple sequences in a batch it was last-writer-
     wins — one sequence's adapter leaked into the others.
  2. Multi-turn: an interactive `ollama run` chat continues one KV cache,
     so turn 2 never saw position 0 and the index never reset — the
     adapter stayed stuck on across turns.

Port the vLLM/HF backend mechanism faithfully: a single-head causal
"router" attention recovers the adapter index in-graph. Per token, only
dim 0 carries signal — Q[0]=1, K[0]=+gain for a control token / -gain
otherwise, V[0]=adapter slot / 0 — and the causal softmax over the single
visible control token recovers that adapter's slot (readback =
clamp(round(V[0]), 0, n_adapters)). gain=15 matches config.py and is
F16-safe (no F32 cache).

The router's K/V live in the model KV cache at an extra layer
R == hparams.router_layer (== n_layer). We bump n_layer_all to n_real+1
so the cache allocator gives the router its own per-sequence slot, and
set n_layer_nextn=1 so n_layer() stays n_real — the decoder loop and
tensor loading are untouched and never reference layer R. The router K is
exempted from the k-shift RoPE loop (its dim-0 value is a literal
magnitude, not a rotation).

Because the selection now lives in the per-sequence KV cache, CONCURRENT
requests are isolated for free (problem 1 fixed; verified by
scratch/concurrent_switch_test.cpp). set_input becomes stateless pure
per-token maps; the global is gone.

Single-switch contract / known limitation, identical to vLLM & HF: the
gain is flat (no recency), so within one sequence there is no mechanism to
revert to base mid-sequence — once an adapter fires it stays on until that
sequence ends (problem 2 is therefore NOT fixed by a faithful copy; vLLM/HF
avoid it only because each served request is a fresh sequence). A client
continuing one KV cache across turns must start a fresh sequence per turn,
or opt into a recency-biased router (a deliberate divergence, not done
here). Documented in granite_switch.cpp and asserted by
scratch/multiturn_leak_test.cpp.

Verified (CPU): both demos unchanged (answerability -> "unanswerable",
query_rewrite -> rewritten query); concurrent two-sequence isolation
passes; multi-turn carry-over matches the vLLM/HF contract.

* granite-switch: drop scratch tests and mac demo for upstream PR

Remove the local-only development artifacts that should not ship in the
upstream PR:
  - granite-switch-mac-demo.sh (local Metal build + demo driver)
  - scratch/concurrent_switch_test.cpp
  - scratch/multiturn_leak_test.cpp

Also drop the now-dangling reference to the scratch tests from the
granite_switch.cpp header comment. Leaves only the core architecture
support (conversion, gguf constants, llama-arch/model/kv-cache, and the
granite_switch graph).

* granite-switch: trim comments to match native llama.cpp style

* granite-switch: trim conversion comments to match native style

* granite-switch: drop unused adapter_ranks metadata

* granite-switch: rename arch to graniteswitch and drop obid alias

* granite-switch: fix non-ASCII comments and document router gain assumption

* granite-switch: drop section comments from constants.py to match native style

* granite-switch: add functional tensor block comments matching Granite4 Vision style

* granite-switch: clarify n_expert_used comment

State the actual constraint: mul_mat_id needs n_expert_used == 1, and
since the GGUF carries expert_count = 0 the generic loader's
n_expert == 0 => n_expert_used == 0 assertion has already passed by the
time load_arch_hparams runs, so it is forced to 1 here.

* granite-switch: note n_layer_nextn reuse has no MTP

The router carving reuses n_layer_nextn, normally the MTP/next-token
count. Clarify in the comment that it is borrowed here purely as the
trailing-layers lever and that there is no MTP head, to spare readers
the double-take.

* granite-switch: rename source file and apply review nits

* granite-switch: don't force LoRA tensors to F16, follow --outtype instead

* granite-switch: drop redundant _permute_qk wrapper, call LlamaModel.permute directly

* granite-switch: read router gain from GGUF (control_token_gain) instead of hardcoding 15.0

* granite-switch: derive n_slots()

* granite-switch: move llm_graph_input_switch into granite-switch.cpp

* granite-switch: cut AI-style narration comments

* granite-switch: collapse multi-line comments

* granite-switch: rename control_token_* maps to adapter_token_*

* granite-switch: cut noise comments

* granite-switch: rename embedded LoRA tensors to <base>.lora_a/lora_b

* granite-switch: GGML_ASSERT token input to avoid UB on embeddings

* granite-switch: TODO for raw embedding input support

* granite-switch: collapse LoRA tensor constants to .lora_a/.lora_b suffix

* granite-switch: drop n_expert_used hack, guard mul_mat_id buft probe

* granite-switch: stop forcing dense expert counts, read from config

* granite-switch: renamed control_token_gain metadata key to router_gain

* granite-switch: trim header comments to match native style

* granite-switch: collapse LoRA tensors to base name + suffix

* granite-switch: inline suffix checks in tensor op resolution

* granite-switch: drop switch-lora struct comment

* granite-switch: guard router layer index and inline n_slots

* granite-switch: group adapter metadata under {arch}.adapters.* namespace

* granite-switch: add hparams.has_rope(il) for KV-shift rope skipping

* granite-switch: skip arch in test-llama-archs (adapter fixture missing, TODO)

* granite-switch: Keys.Adapters namespace + simplify n_slots

* granite-switch: validate substitute token ids against n_vocab

* granite-switch: bound adapter count and lora rank from GGUF

* granite-switch: reject MTP context type when router_layer is set

* granite-switch: throw on bad adapter metadata instead of GGML_ASSERT

* granite-switch: use ASCII +/- in router K signal comment

* granite-switch: document n_layer_nextn repurpose and its leak points

* granite-switch: gate lora_a/lora_b op mapping on router_layer

* granite-switch: label all three preview model sizes
2026-08-10 09:53:46 +02:00

375 lines
14 KiB
Python
Raw Blame History

from __future__ import annotations
from .base import (
ModelBase, TextModel, MmprojModel, ModelType, SentencePieceTokenTypes,
logger, _mistral_common_installed, _mistral_import_error_msg,
get_model_architecture, LazyTorchTensor,
)
from typing import Type
__all__ = [
"ModelBase", "TextModel", "MmprojModel", "ModelType", "SentencePieceTokenTypes",
"get_model_architecture", "LazyTorchTensor", "logger",
"_mistral_common_installed", "_mistral_import_error_msg",
"get_model_class", "print_registered_models", "load_all_models",
]
TEXT_MODEL_MAP: dict[str, str] = {
"AfmoeForCausalLM": "afmoe",
"LagunaForCausalLM": "laguna",
"ApertusForCausalLM": "llama",
"ArceeForCausalLM": "llama",
"ArcticForCausalLM": "arctic",
"AudioFlamingo3ForConditionalGeneration": "qwen",
"BaiChuanForCausalLM": "baichuan",
"BaichuanForCausalLM": "baichuan",
"BailingMoeForCausalLM": "bailingmoe",
"BailingMoeV2ForCausalLM": "bailingmoe",
"BambaForCausalLM": "granite",
"BertForMaskedLM": "bert",
"BertForSequenceClassification": "bert",
"BertModel": "bert",
"BitnetForCausalLM": "bitnet",
"BitNetForCausalLM": "bitnet",
"BloomForCausalLM": "bloom",
"BloomModel": "bloom",
"CamembertModel": "bert",
"ChameleonForCausalLM": "chameleon",
"ChameleonForConditionalGeneration": "chameleon",
"ChatGLMForConditionalGeneration": "chatglm",
"ChatGLMModel": "chatglm",
"CodeShellForCausalLM": "codeshell",
"CogVLMForCausalLM": "cogvlm",
"Cohere2MoeForCausalLM": "command_r",
"Cohere2ForCausalLM": "command_r",
"CohereForCausalLM": "command_r",
"DbrxForCausalLM": "dbrx",
"DeciLMForCausalLM": "deci",
"DeepseekForCausalLM": "deepseek",
"DeepseekOCRForCausalLM": "deepseek",
"DeepseekV2ForCausalLM": "deepseek",
"DeepseekV3ForCausalLM": "deepseek",
"DeepseekV32ForCausalLM": "deepseek",
"DFlashDraftModel": "qwen",
"Qwen3DSparkModel": "qwen",
"DeepseekV4ForCausalLM": "deepseek",
"DeepseekV4DSparkModel": "deepseek",
"DistilBertForMaskedLM": "bert",
"DistilBertForSequenceClassification": "bert",
"DistilBertModel": "bert",
"Dots1ForCausalLM": "dots1",
"DotsOCRForCausalLM": "qwen",
"DreamModel": "dream",
"Ernie4_5ForCausalLM": "ernie",
"Ernie4_5_ForCausalLM": "ernie",
"Ernie4_5_MoeForCausalLM": "ernie",
"EuroBertModel": "bert",
"Exaone4_5_ForConditionalGeneration": "exaone",
"Exaone4ForCausalLM": "exaone",
"ExaoneForCausalLM": "exaone",
"ExaoneMoEForCausalLM": "exaone",
"ExaoneMoeForCausalLM": "exaone",
"FalconForCausalLM": "falcon",
"FalconH1ForCausalLM": "falcon_h1",
"FalconMambaForCausalLM": "mamba",
"GPT2LMHeadModel": "gpt2",
"GPTBigCodeForCausalLM": "starcoder",
"GPTNeoXForCausalLM": "gptneox",
"GPTRefactForCausalLM": "refact",
"Gemma2ForCausalLM": "gemma",
"Gemma3ForCausalLM": "gemma",
"Gemma3ForConditionalGeneration": "gemma",
"Gemma3TextModel": "gemma",
"Gemma3nForCausalLM": "gemma",
"Gemma3nForConditionalGeneration": "gemma",
"Gemma4AssistantForCausalLM": "gemma",
"Gemma4ForConditionalGeneration": "gemma",
"Gemma4ForCausalLM": "gemma",
"Gemma4UnifiedForConditionalGeneration": "gemma",
"Gemma4UnifiedAssistantForCausalLM": "gemma",
"GemmaForCausalLM": "gemma",
"Glm4ForCausalLM": "glm",
"Glm4MoeForCausalLM": "glm",
"Glm4MoeLiteForCausalLM": "glm",
"Glm4vForConditionalGeneration": "glm",
"Glm4vMoeForConditionalGeneration": "glm",
"GlmForCausalLM": "chatglm",
"GlmMoeDsaForCausalLM": "glm",
"GlmOcrForConditionalGeneration": "glm",
"GptOssForCausalLM": "gpt_oss",
"GraniteForCausalLM": "granite",
"GraniteMoeForCausalLM": "granite",
"GraniteMoeHybridForCausalLM": "granite",
"GraniteMoeSharedForCausalLM": "granite",
"GraniteSwitchForCausalLM": "granite",
"GraniteSpeechForConditionalGeneration": "granite",
"GraniteSpeechPlusForConditionalGeneration": "granite",
"Grok1ForCausalLM": "grok",
"GrokForCausalLM": "grok",
"GroveMoeForCausalLM": "grovemoe",
"HunYuanDenseV1ForCausalLM": "hunyuan",
"HunYuanMoEV1ForCausalLM": "hunyuan",
"HunYuanVLForConditionalGeneration": "hunyuan",
"HYV3ForCausalLM": "hunyuan",
"IQuestCoderForCausalLM": "llama",
"InternLM2ForCausalLM": "internlm",
"InternLM3ForCausalLM": "internlm",
"JAISLMHeadModel": "jais",
"Jais2ForCausalLM": "jais",
"JambaForCausalLM": "jamba",
"JanusForConditionalGeneration": "januspro",
"JinaBertForMaskedLM": "bert",
"JinaBertModel": "bert",
"JinaEmbeddingsV5Model": "bert",
"KORMoForCausalLM": "qwen",
"KimiK25ForConditionalGeneration": "deepseek",
"KimiLinearForCausalLM": "kimi_linear",
"KimiLinearModel": "kimi_linear",
"KimiVLForConditionalGeneration": "deepseek",
"LFM2ForCausalLM": "lfm2",
"LLaDAMoEModel": "llada",
"LLaDAMoEModelLM": "llada",
"LLaDAModelLM": "llada",
"LLaMAForCausalLM": "llama",
"Lfm25AudioTokenizer": "lfm2",
"Lfm2BidirectionalModel": "lfm2",
"Lfm2ForCausalLM": "lfm2",
"Lfm2Model": "lfm2",
"Lfm2MoeForCausalLM": "lfm2",
"Llama4ForCausalLM": "llama",
"Llama4ForConditionalGeneration": "llama",
"LlamaBidirectionalModel": "llama",
"LlamaForCausalLM": "llama",
"LlamaModel": "llama",
"Eagle3DraftModel": "llama",
"Eagle3Speculator": "llama",
"Eagle3LlamaForCausalLM": "llama",
"LlamaForCausalLMEagle3": "llama",
"LlavaForConditionalGeneration": "llama",
"LlavaStableLMEpochForCausalLM": "stablelm",
"MPTForCausalLM": "mpt",
"MT5ForConditionalGeneration": "t5",
"MaincoderForCausalLM": "maincoder",
"Mamba2ForCausalLM": "mamba",
"MambaForCausalLM": "mamba",
"MambaLMHeadModel": "mamba",
"MellumForCausalLM": "mellum",
"MiMoV2FlashForCausalLM": "mimo",
"MiMoV2ForCausalLM": "mimo",
"MiniCPM3ForCausalLM": "minicpm",
"MiniCPMForCausalLM": "minicpm",
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
"MiniMaxM2ForCausalLM": "minimax",
"MiniMaxM3SparseForCausalLM": "minimax",
"MiniMaxM3SparseForConditionalGeneration": "minimax",
"Ministral3ForCausalLM": "mistral3",
"Mistral3ForConditionalGeneration": "mistral3",
"MistralForCausalLM": "llama",
"MixtralForCausalLM": "llama",
"ModernBertForMaskedLM": "bert",
"ModernBertForSequenceClassification": "bert",
"ModernBertModel": "bert",
"NanbeigeForCausalLM": "nanbeige",
"NemotronForCausalLM": "nemotron",
"NemotronHForCausalLM": "nemotron",
"NeoBERT": "bert",
"NeoBERTForSequenceClassification": "bert",
"NeoBERTLMHead": "bert",
"NomicBertModel": "bert",
"OLMoForCausalLM": "olmo",
"Olmo2ForCausalLM": "olmo",
"Olmo3ForCausalLM": "olmo",
"OlmoForCausalLM": "olmo",
"OlmoeForCausalLM": "olmo",
"OpenELMForCausalLM": "openelm",
"OrionForCausalLM": "orion",
"PLMForCausalLM": "plm",
"PLaMo2ForCausalLM": "plamo",
"PLaMo3ForCausalLM": "plamo",
"PaddleOCRVLForConditionalGeneration": "ernie",
"PanguEmbeddedForCausalLM": "pangu",
"Phi3ForCausalLM": "phi",
"Phi4ForCausalLMV": "phi",
"PhiForCausalLM": "phi",
"PhiMoEForCausalLM": "phi",
"Plamo2ForCausalLM": "plamo",
"Plamo3ForCausalLM": "plamo",
"PlamoForCausalLM": "plamo",
"QWenLMHeadModel": "qwen",
"Qwen2AudioForConditionalGeneration": "qwen",
"Qwen2ForCausalLM": "qwen",
"Qwen2Model": "qwen",
"Qwen2MoeForCausalLM": "qwen",
"Qwen2VLForConditionalGeneration": "qwenvl",
"Qwen2VLModel": "qwenvl",
"Qwen2_5OmniModel": "qwenvl",
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
"Qwen3ASRForConditionalGeneration": "qwen3vl",
"Qwen3ForCausalLM": "qwen",
"Qwen3Model": "qwen",
"Qwen3MoeForCausalLM": "qwen",
"Qwen3NextForCausalLM": "qwen",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
"Qwen3_5ForCausalLM": "qwen",
"Qwen3_5ForConditionalGeneration": "qwen",
"Qwen3_5MoeForCausalLM": "qwen",
"Qwen3_5MoeForConditionalGeneration": "qwen",
"RND1": "qwen",
"RWForCausalLM": "falcon",
"RWKV6Qwen2ForCausalLM": "rwkv",
"RWKV7ForCausalLM": "rwkv",
"RobertaForSequenceClassification": "bert",
"RobertaModel": "bert",
"RuGPT3XLForCausalLM": "gpt2",
"Rwkv6ForCausalLM": "rwkv",
"Rwkv7ForCausalLM": "rwkv",
"RwkvHybridForCausalLM": "rwkv",
"Sarashina2VisionForCausalLM": "sarashina2",
"SarvamMoEForCausalLM": "bailingmoe",
"SeedOssForCausalLM": "olmo",
"SmallThinkerForCausalLM": "smallthinker",
"SmolLM3ForCausalLM": "llama",
"SolarOpenForCausalLM": "glm",
"StableLMEpochForCausalLM": "stablelm",
"StableLmForCausalLM": "stablelm",
"Starcoder2ForCausalLM": "starcoder",
"Step3p5ForCausalLM": "step3",
"StepVLForConditionalGeneration": "step3",
"Step3p7ForConditionalGeneration": "step3",
"T5EncoderModel": "t5",
"T5ForConditionalGeneration": "t5",
"T5WithLMHeadModel": "t5",
"TalkieForCausalLM": "talkie",
"UMT5ForConditionalGeneration": "t5",
"UMT5Model": "t5",
"UltravoxModel": "ultravox",
"UnlimitedOCRForCausalLM": "deepseek",
"VLlama3ForCausalLM": "llama",
"VoxtralForConditionalGeneration": "llama",
"WavTokenizerDec": "wavtokenizer",
"XLMRobertaForSequenceClassification": "bert",
"XLMRobertaModel": "bert",
"XverseForCausalLM": "xverse",
"YoutuForCausalLM": "deepseek",
"YoutuVLForConditionalGeneration": "deepseek",
"modeling_grove_moe.GroveMoeForCausalLM": "grovemoe",
"modeling_sarvam_moe.SarvamMoEForCausalLM": "bailingmoe",
}
MMPROJ_MODEL_MAP: dict[str, str] = {
"AudioFlamingo3ForConditionalGeneration": "ultravox",
"CogVLMForCausalLM": "cogvlm",
"DeepseekOCR2ForCausalLM": "deepseek",
"DeepseekOCRForCausalLM": "deepseek",
"DotsOCRForCausalLM": "dotsocr",
"Exaone4_5_ForConditionalGeneration": "exaone",
"Gemma3ForConditionalGeneration": "gemma",
"Gemma3nForConditionalGeneration": "gemma",
"Gemma4ForConditionalGeneration": "gemma",
"Gemma4UnifiedForConditionalGeneration": "gemma",
"Glm4vForConditionalGeneration": "qwen3vl",
"Glm4vMoeForConditionalGeneration": "qwen3vl",
"Glm5vForConditionalGeneration": "kimivl",
"GlmOcrForConditionalGeneration": "qwen3vl",
"GlmasrModel": "ultravox",
"Granite4VisionForConditionalGeneration": "granite",
"GraniteSpeechForConditionalGeneration": "granite",
"GraniteSpeechPlusForConditionalGeneration": "granite",
"HunYuanVLForConditionalGeneration": "hunyuan",
"Idefics3ForConditionalGeneration": "smolvlm",
"InternVisionModel": "internvl",
"JanusForConditionalGeneration": "januspro",
"KimiK25ForConditionalGeneration": "kimivl",
"KimiVLForConditionalGeneration": "kimivl",
"Lfm2AudioForConditionalGeneration": "lfm2",
"Lfm2VlForConditionalGeneration": "lfm2",
"LightOnOCRForConditionalGeneration": "lighton_ocr",
"Llama4ForConditionalGeneration": "llama4",
"LlavaForConditionalGeneration": "llava",
"MERaLiON2ForConditionalGeneration": "ultravox",
"MiMoV2ForCausalLM": "mimo",
"MiniMaxM3SparseForConditionalGeneration": "minimax",
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
"Mistral3ForConditionalGeneration": "llava",
"NemotronH_Nano_VL_V2": "nemotron",
"PaddleOCRVisionModel": "ernie",
"Phi4ForCausalLMV": "phi",
"Qwen2AudioForConditionalGeneration": "ultravox",
"Qwen2VLForConditionalGeneration": "qwenvl",
"Qwen2VLModel": "qwenvl",
"Qwen2_5OmniModel": "qwenvl",
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
"Qwen3ASRForConditionalGeneration": "qwen3vl",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
"Qwen3_5ForConditionalGeneration": "qwen3vl",
"Qwen3_5MoeForConditionalGeneration": "qwen3vl",
"RADIOModel": "nemotron",
"Sarashina2VisionForCausalLM": "sarashina2",
"SmolVLMForConditionalGeneration": "smolvlm",
"StepVLForConditionalGeneration": "step3",
"Step3p7ForConditionalGeneration": "step3",
"UltravoxModel": "ultravox",
"UnlimitedOCRForCausalLM": "deepseek",
"VoxtralForConditionalGeneration": "ultravox",
"YoutuVLForConditionalGeneration": "youtuvl",
}
_TEXT_MODEL_MODULES = sorted(set(TEXT_MODEL_MAP.values()))
_MMPROJ_MODEL_MODULES = sorted(set(MMPROJ_MODEL_MAP.values()))
_loaded_text_modules: set[str] = set()
_loaded_mmproj_modules: set[str] = set()
def load_all_models() -> None:
"""Import all model modules to trigger @ModelBase.register() decorators."""
if len(_loaded_text_modules) != len(_TEXT_MODEL_MODULES):
for module_name in _TEXT_MODEL_MODULES:
if module_name not in _loaded_text_modules:
try:
__import__(f"conversion.{module_name}")
_loaded_text_modules.add(module_name)
except Exception as e:
logger.warning(f"Failed to load model module {module_name}: {e}")
if len(_loaded_mmproj_modules) != len(_MMPROJ_MODEL_MODULES):
for module_name in _MMPROJ_MODEL_MODULES:
if module_name not in _loaded_mmproj_modules:
try:
__import__(f"conversion.{module_name}")
_loaded_mmproj_modules.add(module_name)
except Exception as e:
logger.warning(f"Failed to load model module {module_name}: {e}")
def get_model_class(name: str, mmproj: bool = False) -> Type[ModelBase]:
"""Dynamically import and return a model class by its HuggingFace architecture name."""
relevant_map = MMPROJ_MODEL_MAP if mmproj else TEXT_MODEL_MAP
if name not in relevant_map:
raise NotImplementedError(f"Architecture {name!r} not supported!")
module_name = relevant_map[name]
__import__(f"conversion.{module_name}")
model_type = ModelType.MMPROJ if mmproj else ModelType.TEXT
return ModelBase._model_classes[model_type][name]
def print_registered_models() -> None:
load_all_models()
logger.error("TEXT models:")
for name in sorted(TEXT_MODEL_MAP.keys()):
logger.error(f" - {name}")
logger.error("MMPROJ models:")
for name in sorted(MMPROJ_MODEL_MAP.keys()):
logger.error(f" - {name}")
Reference in New Issue View Git Blame Copy Permalink