4625240437
* model: K2 Horizon gguf conversion code
* model: loading hparams and tensors in k2-horizon.cpp
* model: K2 Horizon compute graph
* model: K2 Horizon compute graph adjustment and registering tokenizers
* model: K2 Horizon chat template and accomodate safetensors naming
* unicode : add the K2-Horizon pre-tokenizer splitter
The K2-Horizon regex had no arm in unicode_regex_split_custom and fell through to the
general std::regex fallback, which fails two ways.
On MSVC std::regex rejects \p{...}, so no K2-Horizon GGUF loads on Windows at all:
llama-quantize, llama-imatrix and llama-perplexity all abort with
regex_error(error_escape) before a token is produced.
Where the fallback does compile it is still wrong. unicode_regex_split collapses each
codepoint to a single byte naming its Unicode category before matching, and U+200C/U+200D
are category Control, which has no entry in k_ucat_cpt, so both become the 0xD0 fallback
byte. The literal and alternatives in K2's regex can then never match and
every ZWNJ or ZWJ ends a letter run.
The splitter is the existing llama3 one with a single rule widened, since K2's regex
differs from llama3's only in that a letter run also takes marks, ZWNJ and ZWJ.
tests/test-unicode.cpp gains a case for this: it fails before the change with
[Amy] [ZWNJ khaham] and passes after with the run intact.
* tests: expand K2 Horizon unicode splitter coverage
* unicode: handle K2 Horizon case folding and empty input
Assisted-by: Codex
* jinja : support sequence indices in selectattr and rejectattr
Assisted-by: Codex
* model : add K2 Horizon dense and MoVA support
Includes the K2 Horizon implementation from ifm-ai/llama.cpp with converter, tensor-parallel and model save/reload fixes.
Assisted-by: Codex
* chat : support K2 Horizon reasoning and tool calls
Assisted-by: Codex
* conversion: remove obsolete K2 Aurora alias
Assisted-by: Codex
* k2-horizon: enforce response schemas and load YaRN betas
Constrain final JSON after reasoning, accept flexible JSON tool envelopes,
enforce XML dialects, and handle repeated or alternate thinking markers.
Load YaRN beta metadata instead of retaining the default values.
Add schema, streaming, continuation, and model reload regressions. Validate
CUDA and CPU builds and 0.9B, 4B, and MoVA conversation/tool round trips.
Assisted-by: Codex
* renaming template fixture
* adressing cisc follows ups
* desloppify the parser / adress aldehir comments
* clean test-chat
* remove fallback : model trained mostly on high anyway
* fix k2 attn_v_exp tn splitting and metal fusion baseline
* k2-horizon : forward expand views before sums
* k2-horizon: copy embds before group norm to fix TP
* disable tesnor parallelism
---------
Co-authored-by: Ryandito Diandaru <ryandito.diandaru@mbzuai.ac.ae>
Co-authored-by: WestWaters <mario.papaleo2013@gmail.com>
Co-authored-by: Natani L. Mayday <71436458+TaskPuppyNatani@users.noreply.github.com>
Co-authored-by: West <100190545+WestWaters@users.noreply.github.com>
Co-authored-by: aaryamonvikram <aaryamonvikram@gmail.com>
Co-authored-by: aaryamonvikram <96529820+aaryamonvikram@users.noreply.github.com>
424 lines
16 KiB
Python
424 lines
16 KiB
Python
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",
|
|
"BailingMoeV3VLForConditionalGeneration": "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",
|
|
"ClefModel": "clef",
|
|
"CodeShellForCausalLM": "codeshell",
|
|
"CogVLMForCausalLM": "cogvlm",
|
|
"Cohere2MoeForCausalLM": "command_r",
|
|
"Cohere2ForCausalLM": "command_r",
|
|
"CohereForCausalLM": "command_r",
|
|
"DbrxForCausalLM": "dbrx",
|
|
"DeciLMForCausalLM": "deci",
|
|
"PplxDeciderModel": "pplx_decider",
|
|
"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",
|
|
"EmbeddingGemma2Model": "gemma",
|
|
"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",
|
|
"Gemma4DSparkModel": "gemma",
|
|
"Gemma4ForConditionalGeneration": "gemma",
|
|
"Gemma4ForCausalLM": "gemma",
|
|
"Gemma4UnifiedForConditionalGeneration": "gemma",
|
|
"Gemma4UnifiedAssistantForCausalLM": "gemma",
|
|
"GemmaForCausalLM": "gemma",
|
|
"Glm4ForCausalLM": "glm",
|
|
"Glm4MoeForCausalLM": "glm",
|
|
"Glm4MoeLiteForCausalLM": "glm",
|
|
"Glm4vForConditionalGeneration": "glm",
|
|
"Glm4vMoeForConditionalGeneration": "glm",
|
|
"Glm5NextForConditionalGeneration": "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",
|
|
"HrmTextForCausalLM": "hrm_text",
|
|
"HYV3ForCausalLM": "hunyuan",
|
|
"HYV4ForCausalLM": "hy_v4",
|
|
"IQuestCoderForCausalLM": "llama",
|
|
"InternLM2ForCausalLM": "internlm",
|
|
"InternLM3ForCausalLM": "internlm",
|
|
"JAISLMHeadModel": "jais",
|
|
"Jais2ForCausalLM": "jais",
|
|
"JambaForCausalLM": "jamba",
|
|
"JanusForConditionalGeneration": "januspro",
|
|
"JinaBertForMaskedLM": "bert",
|
|
"JinaBertModel": "bert",
|
|
"JinaEmbeddingsV5Model": "bert",
|
|
"K2HorizonForCausalLM": "k2_horizon",
|
|
"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",
|
|
"KevModel": "lev",
|
|
"LevModel": "lev",
|
|
"NimbleModel": "lev",
|
|
"Lfm25AudioTokenizer": "lfm2",
|
|
"Lfm2BidirectionalForMaskedLM": "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",
|
|
"MapleForCausalLM": "maple",
|
|
"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",
|
|
"ModernBertDecisionModel": "bert",
|
|
"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",
|
|
"OpenJevModel": "qwen",
|
|
"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",
|
|
"Spark2_5ForCausalLM": "spark2_5",
|
|
"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",
|
|
"ClefModel": "clef",
|
|
"CogVLMForCausalLM": "cogvlm",
|
|
"PplxDeciderModel": "pplx_decider",
|
|
"DeepseekOCR2ForCausalLM": "deepseek",
|
|
"DeepseekOCRForCausalLM": "deepseek",
|
|
"DeepseekV4ForCausalLM": "deepseek",
|
|
"Dots3NoteForCausalLM": "dots3",
|
|
"Dots3NoteForConditionalGeneration": "dots3",
|
|
"DotsOCRForCausalLM": "dotsocr",
|
|
"EmbeddingGemma2Model": "gemma",
|
|
"Exaone4_5_ForConditionalGeneration": "exaone",
|
|
"Gemma3ForConditionalGeneration": "gemma",
|
|
"Gemma3nForConditionalGeneration": "gemma",
|
|
"Gemma4ForConditionalGeneration": "gemma",
|
|
"Gemma4UnifiedForConditionalGeneration": "gemma",
|
|
"Glm4vForConditionalGeneration": "qwen3vl",
|
|
"BailingMoeV3VLForConditionalGeneration": "bailingmoe3",
|
|
"Glm4vMoeForConditionalGeneration": "qwen3vl",
|
|
"Glm5NextForConditionalGeneration": "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",
|
|
"OpenJevModel": "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}")
|