Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-04 03:51:05 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
95ef7fc16054e63b427a3ef00188e055ef7586d8
llama.cpp/conversion/__init__.py
T
Yaniss Amazouz c61b98b875 model: add NVIDIA Nemotron-3-Puzzle-75B-A9B (NemotronHPuzzle) support (#25444) 
* hparams: add per-layer n_ff_exp/n_expert_used arrays with scalar-or-array loading

G1/G2 infrastructure for variable-per-layer expert FFN size and top-k routing
(required for Puzzle-75B which has 5 distinct n_ff_exp values and 7 top-k values
across its 40 MoE layers).

Design: rename scalar members to _impl suffix (following existing convention),
add LLAMA_MAX_LAYERS arrays, add n_ff_exp(il)/n_expert_used(il) accessors with
scalar fallback. No new GGUF keys: reuses existing expert_feed_forward_length and
expert_used_count keys via get_key_or_arr (scalar -> broadcast, array -> per-layer).

- llama-hparams.h: n_ff_exp -> n_ff_exp_impl, n_expert_used -> n_expert_used_impl;
  add n_ff_exp_arr / n_expert_used_arr arrays; add per-layer accessor declarations.
- llama-hparams.cpp: implement n_ff_exp(il) and n_expert_used(il); out-of-range
  il returns impl safely (shared code, no abort).
- llama-model.cpp: central n_expert_used load changed to get_key_or_arr; derive
  impl as max-of-array for validations and backward compat; zero both new arrays;
  HunyuanVL override also zeroes n_expert_used_arr.
- llama-graph.cpp: aggregation loop in build_moe_ffn uses hparams.n_expert_used(il)
  so per-layer top-k bounds the ggml_view loop correctly.
- All other files: mechanical rename hparams.n_{ff_exp,expert_used} -> *_impl.
  Scalar arches are unaffected (broadcast fills all array slots with the single value).

(cherry picked from commit 269a81e03d66e1c353e1a203a0c03a03eb2b1a4e)

* nemotron-h: use per-layer n_ff_exp(il) and n_expert_used(il) at MoE call-sites

Load n_ff_exp via get_key_or_arr into hparams.n_ff_exp_arr in load_arch_hparams;
derive impl as max for existing uniform GGUFs.

In load_arch_tensors, compute n_ff_exp_i = hparams.n_ff_exp(i) with fallback to
n_ff(i)/n_expert_used(i) for GGUFs that omit expert_feed_forward_length.

In build_ffn_layer, pass hparams.n_expert_used(il) to build_moe_ffn so per-layer
top-k is used for expert routing selection.

All other nemotron-h behaviour (mamba2, attention, shared-exp, latent projection,
routed_scaling_factor, expert_weights_norm, sigmoid gating) is unchanged.

(cherry picked from commit b1878a101793cd4e59868ac72635a86ea694987c)

* arch/*.cpp + gguf-py: mechanical rename n_ff_exp->n_ff_exp_impl, n_expert_used->n_expert_used_impl

All non-nemotron arch files continue using the scalar impl member directly.
Behaviour is identical: the impl value is the broadcast value from the GGUF scalar.

gguf_writer: add_expert_feed_forward_length and add_expert_used_count now accept
int | Sequence[int], mirroring add_feed_forward_length, so converters can write
per-layer arrays with the same existing GGUF keys.

(cherry picked from commit 8f009f54bea5ef9a6a354123bd25e9d5ea2d5e03)

* convert: support NemotronHPuzzleForCausalLM (per-block MoE config)

Parse block_configs/mtp_block_configs into per-layer arrays (scalar-or-array
keys), append the MTP [attention, moe] sub-blocks as blk.88/blk.89 with
nextn tensors, accept the backbone.* prefix, and register the arch.
Also fix a pre-existing undeclared _experts attribute on NemotronHModel.

(cherry picked from commit d1a592f278336e78457454eb6c96bca917135f10)

* nemotron-h: distinguish Nemotron 3 Puzzle (75B.A9B) from Super (120B.A12B)

Both have 88 layers; the per-layer expert_used_count array (heterogeneous
for Puzzle, broadcast-uniform for Super) is the discriminator.

(cherry picked from commit f824e09dc812169589cf5662c92d149a4c18c30a)

* convert: accept the official Puzzle BF16 checkpoint's tensor naming

The officially distributed BF16 checkpoint (NVIDIA-Nemotron-Labs-3-Puzzle-
75B-A9B-BF16) names the trunk model.* (model.layers.*, model.embeddings,
model.norm_f) where the original release used the NemotronH-style
backbone.*, and spells the router bias e_score_correction_bias instead of
e_score_correction.bias. Normalize both at the top of
NemotronHPuzzleModel.modify_tensors so either checkpoint converts; every
tensor name in the official index (42683 keys, MTP head included) resolves
through the tensor map after normalization.

(cherry picked from commit 189b67fc2c9d50970416c94b3317a6e7baa49b03)

* laguna: use n_ff_exp_impl for the uniform-MoE FFN size

Laguna landed after this branch was cut and reads hparams.n_ff_exp as a
scalar. This series turns it into a per-layer array with an n_ff_exp(il)
accessor, so the three scalar reads no longer compile. Laguna is a
uniform MoE, so point them at the scalar fallback n_ff_exp_impl, same as
deepseek2/qwen3moe/gemma4 in this series. No behaviour change.

(cherry picked from commit dbedc9e19c50dca0acdfb402362e2707bee424ae)

* arch: extend the n_ff_exp/n_expert_used rename to archs added upstream

kimi-k3, dflash, bailingmoe3, deepseek4, granite-swa and the nemotron-h MTP
block still referenced the scalar fields by their old names. n_ff_exp and
n_expert_used are accessors now, so those reads no longer compile; point the
non-per-layer archs at the _impl scalars and use the indexed form where the
call site is per-layer.

* convert: keep Puzzle opted out of the NemotronH MTP export path

#26725 added MTP export to NemotronHModel, keyed on num_nextn_predict_layers.
Puzzle's config carries that key, but NemotronHPuzzleModel bypasses
NemotronHModel.__init__ (its per-block config needs a different setup), so
_mtp_bid was never assigned and modify_tensors raised AttributeError on any
mtp.* tensor. Puzzle's head is also laid out by mtp_block_configs, not the
mtp.layers.* form the base maps.

Set _mtp_bid to None, drop mtp.* in filter_tensors, and declare
supports_mtp_export = False so --mtp / --no-mtp fail at the CLI.

* llama: replace n_ff_exp/n_expert_used scalars with per-layer accessors

Follow-up to review feedback: the previous revision kept the scalar
hparams fields alongside the new per-layer arrays, which duplicated
state that get_key_or_arr already handles by broadcasting a scalar
value over every layer.

Drop both scalars and expose n_ff_exp(il) / n_expert_used(il) built
exactly like the existing n_head_kv(il) and n_ff(il) accessors: they
index the array and GGML_ABORT out of range, with il defaulting to 0
so genuinely uniform call sites stay a plain n_ff_exp().

Arch loaders now read both keys through get_key_or_arr over
n_layer_all, and the n_expert_used validation checks the maximum
across layers instead of a single field.

* llama: restore per-key required flags on the expert hparam reads

The scalar-to-array conversion passed required=false at every call site,
which silently made mandatory keys optional. Each read now carries the
same required flag it had before the conversion.
2026-09-03 08:53:08 +02:00

401 lines
15 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",
"BailingMoeV3ForCausalLM": "bailingmoe3",
"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",
"DFlash2DraftModel": "qwen",
"Qwen3DSparkModel": "qwen",
"DSparkDraftModel": "qwen",
"DSparkSpeculator": "qwen",
"Lfm2DSparkDraftModel": "qwen",
"LingDSparkModel": "qwen",
"DeepseekV4ForCausalLM": "deepseek",
"DeepseekV4DSparkModel": "deepseek",
"DistilBertForMaskedLM": "bert",
"DistilBertForSequenceClassification": "bert",
"DistilBertModel": "bert",
"Dots1ForCausalLM": "dots1",
"Dots3NoteForCausalLM": "dots3",
"Dots3NoteForConditionalGeneration": "dots3",
"Dots3NoteTextForCausalLM": "dots3",
"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",
"GraniteSWAForCausalLM": "granite",
"GraniteMoeSWAForCausalLM": "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",
"KimiK3ForConditionalGeneration": "kimi_k3",
"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",
"MiniMaxText01ForCausalLM": "minimax",
"MiniMaxM1ForCausalLM": "minimax",
"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",
"NemotronHPuzzleForCausalLM": "nemotron",
"NeoBERT": "bert",
"NeoBERTForSequenceClassification": "bert",
"NeoBERTLMHead": "bert",
"NomicBertModel": "bert",
"OLMoForCausalLM": "olmo",
"Olmo2ForCausalLM": "olmo",
"Olmo3ForCausalLM": "olmo",
"OlmoForCausalLM": "olmo",
"OlmoeForCausalLM": "olmo",
"MuseGlimmerAssistantModel": "muse_glimmer",
"MuseGlimmerForConditionalGeneration": "muse_glimmer",
"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",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
"Qwen3_5ForCausalLM": "qwen",
"Qwen3_5ForConditionalGeneration": "qwen",
"Qwen3_5MoeForCausalLM": "qwen",
"Qwen3_5MoeForConditionalGeneration": "qwen",
"Qwen4ExpForCausalLM": "qwen4exp",
"Qwen4ExpForConditionalGeneration": "qwen4exp",
"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",
"DeepseekV4ForCausalLM": "deepseek",
"Dots3NoteForCausalLM": "dots3",
"Dots3NoteForConditionalGeneration": "dots3",
"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",
"MuseGlimmerForConditionalGeneration": "muse_glimmer",
"PaddleOCRVisionModel": "ernie",
"Phi4ForCausalLMV": "phi",
"Qwen2AudioForConditionalGeneration": "ultravox",
"Qwen2VLForConditionalGeneration": "qwenvl",
"Qwen2VLModel": "qwenvl",
"Qwen2_5OmniModel": "qwenvl",
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
"Qwen3ASRForConditionalGeneration": "qwen3vl",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
"Qwen3_5ForConditionalGeneration": "qwen3vl",
"Qwen3_5MoeForConditionalGeneration": "qwen3vl",
"Qwen4ExpForConditionalGeneration": "qwen4exp",
"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