char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Alexllama: add Maple 20B-A1B ternary MoE architecture (CPU) (#27000)3d10bcd19

char/aliceai-t5-moe
15.5 KiB403 linesraw
1from __future__ import annotations
2
3from .base import (
4    ModelBase, TextModel, MmprojModel, ModelType, SentencePieceTokenTypes,
5    logger, _mistral_common_installed, _mistral_import_error_msg,
6    get_model_architecture, LazyTorchTensor,
7)
8from typing import Type
9
10
11__all__ = [
12    "ModelBase", "TextModel", "MmprojModel", "ModelType", "SentencePieceTokenTypes",
13    "get_model_architecture", "LazyTorchTensor", "logger",
14    "_mistral_common_installed", "_mistral_import_error_msg",
15    "get_model_class", "print_registered_models", "load_all_models",
16]
17
18
19TEXT_MODEL_MAP: dict[str, str] = {
20    "AfmoeForCausalLM": "afmoe",
21    "LagunaForCausalLM": "laguna",
22    "ApertusForCausalLM": "llama",
23    "ArceeForCausalLM": "llama",
24    "ArcticForCausalLM": "arctic",
25    "AudioFlamingo3ForConditionalGeneration": "qwen",
26    "BaiChuanForCausalLM": "baichuan",
27    "BaichuanForCausalLM": "baichuan",
28    "BailingMoeForCausalLM": "bailingmoe",
29    "BailingMoeV2ForCausalLM": "bailingmoe",
30    "BailingMoeV3ForCausalLM": "bailingmoe3",
31    "BambaForCausalLM": "granite",
32    "BertForMaskedLM": "bert",
33    "BertForSequenceClassification": "bert",
34    "BertModel": "bert",
35    "BitnetForCausalLM": "bitnet",
36    "BitNetForCausalLM": "bitnet",
37    "BloomForCausalLM": "bloom",
38    "BloomModel": "bloom",
39    "CamembertModel": "bert",
40    "ChameleonForCausalLM": "chameleon",
41    "ChameleonForConditionalGeneration": "chameleon",
42    "ChatGLMForConditionalGeneration": "chatglm",
43    "ChatGLMModel": "chatglm",
44    "CodeShellForCausalLM": "codeshell",
45    "CogVLMForCausalLM": "cogvlm",
46    "Cohere2MoeForCausalLM": "command_r",
47    "Cohere2ForCausalLM": "command_r",
48    "CohereForCausalLM": "command_r",
49    "DbrxForCausalLM": "dbrx",
50    "DeciLMForCausalLM": "deci",
51    "DeepseekForCausalLM": "deepseek",
52    "DeepseekOCRForCausalLM": "deepseek",
53    "DeepseekV2ForCausalLM": "deepseek",
54    "DeepseekV3ForCausalLM": "deepseek",
55    "DeepseekV32ForCausalLM": "deepseek",
56    "DFlashDraftModel": "qwen",
57    "DFlash2DraftModel": "qwen",
58    "Qwen3DSparkModel": "qwen",
59    "DSparkDraftModel": "qwen",
60    "DSparkSpeculator": "qwen",
61    "Lfm2DSparkDraftModel": "qwen",
62    "LingDSparkModel": "qwen",
63    "DeepseekV4ForCausalLM": "deepseek",
64    "DeepseekV4DSparkModel": "deepseek",
65    "DistilBertForMaskedLM": "bert",
66    "DistilBertForSequenceClassification": "bert",
67    "DistilBertModel": "bert",
68    "Dots1ForCausalLM": "dots1",
69    "Dots3NoteForCausalLM": "dots3",
70    "Dots3NoteForConditionalGeneration": "dots3",
71    "Dots3NoteTextForCausalLM": "dots3",
72    "DotsOCRForCausalLM": "qwen",
73    "DreamModel": "dream",
74    "Ernie4_5ForCausalLM": "ernie",
75    "Ernie4_5_ForCausalLM": "ernie",
76    "Ernie4_5_MoeForCausalLM": "ernie",
77    "EuroBertModel": "bert",
78    "Exaone4_5_ForConditionalGeneration": "exaone",
79    "Exaone4ForCausalLM": "exaone",
80    "ExaoneForCausalLM": "exaone",
81    "ExaoneMoEForCausalLM": "exaone",
82    "ExaoneMoeForCausalLM": "exaone",
83    "FalconForCausalLM": "falcon",
84    "FalconH1ForCausalLM": "falcon_h1",
85    "FalconMambaForCausalLM": "mamba",
86    "GPT2LMHeadModel": "gpt2",
87    "GPTBigCodeForCausalLM": "starcoder",
88    "GPTNeoXForCausalLM": "gptneox",
89    "GPTRefactForCausalLM": "refact",
90    "Gemma2ForCausalLM": "gemma",
91    "Gemma3ForCausalLM": "gemma",
92    "Gemma3ForConditionalGeneration": "gemma",
93    "Gemma3TextModel": "gemma",
94    "Gemma3nForCausalLM": "gemma",
95    "Gemma3nForConditionalGeneration": "gemma",
96    "Gemma4AssistantForCausalLM": "gemma",
97    "Gemma4ForConditionalGeneration": "gemma",
98    "Gemma4ForCausalLM": "gemma",
99    "Gemma4UnifiedForConditionalGeneration": "gemma",
100    "Gemma4UnifiedAssistantForCausalLM": "gemma",
101    "GemmaForCausalLM": "gemma",
102    "Glm4ForCausalLM": "glm",
103    "Glm4MoeForCausalLM": "glm",
104    "Glm4MoeLiteForCausalLM": "glm",
105    "Glm4vForConditionalGeneration": "glm",
106    "Glm4vMoeForConditionalGeneration": "glm",
107    "GlmForCausalLM": "chatglm",
108    "GlmMoeDsaForCausalLM": "glm",
109    "GlmOcrForConditionalGeneration": "glm",
110    "GptOssForCausalLM": "gpt_oss",
111    "GraniteForCausalLM": "granite",
112    "GraniteMoeForCausalLM": "granite",
113    "GraniteMoeHybridForCausalLM": "granite",
114    "GraniteMoeSharedForCausalLM": "granite",
115    "GraniteSwitchForCausalLM": "granite",
116    "GraniteSpeechForConditionalGeneration": "granite",
117    "GraniteSpeechPlusForConditionalGeneration": "granite",
118    "GraniteSWAForCausalLM": "granite",
119    "GraniteMoeSWAForCausalLM": "granite",
120    "Grok1ForCausalLM": "grok",
121    "GrokForCausalLM": "grok",
122    "GroveMoeForCausalLM": "grovemoe",
123    "HunYuanDenseV1ForCausalLM": "hunyuan",
124    "HunYuanMoEV1ForCausalLM": "hunyuan",
125    "HunYuanVLForConditionalGeneration": "hunyuan",
126    "HYV3ForCausalLM": "hunyuan",
127    "HYV4ForCausalLM": "hy_v4",
128    "IQuestCoderForCausalLM": "llama",
129    "InternLM2ForCausalLM": "internlm",
130    "InternLM3ForCausalLM": "internlm",
131    "JAISLMHeadModel": "jais",
132    "Jais2ForCausalLM": "jais",
133    "JambaForCausalLM": "jamba",
134    "JanusForConditionalGeneration": "januspro",
135    "JinaBertForMaskedLM": "bert",
136    "JinaBertModel": "bert",
137    "JinaEmbeddingsV5Model": "bert",
138    "KORMoForCausalLM": "qwen",
139    "KimiK25ForConditionalGeneration": "deepseek",
140    "KimiK3ForConditionalGeneration": "kimi_k3",
141    "KimiLinearForCausalLM": "kimi_linear",
142    "KimiLinearModel": "kimi_linear",
143    "KimiVLForConditionalGeneration": "deepseek",
144    "LFM2ForCausalLM": "lfm2",
145    "LLaDAMoEModel": "llada",
146    "LLaDAMoEModelLM": "llada",
147    "LLaDAModelLM": "llada",
148    "LLaMAForCausalLM": "llama",
149    "Lfm25AudioTokenizer": "lfm2",
150    "Lfm2BidirectionalModel": "lfm2",
151    "Lfm2ForCausalLM": "lfm2",
152    "Lfm2Model": "lfm2",
153    "Lfm2MoeForCausalLM": "lfm2",
154    "Llama4ForCausalLM": "llama",
155    "Llama4ForConditionalGeneration": "llama",
156    "LlamaBidirectionalModel": "llama",
157    "LlamaForCausalLM": "llama",
158    "LlamaModel": "llama",
159    "Eagle3DraftModel": "llama",
160    "Eagle3Speculator": "llama",
161    "Eagle3LlamaForCausalLM": "llama",
162    "LlamaForCausalLMEagle3": "llama",
163    "LlavaForConditionalGeneration": "llama",
164    "LlavaStableLMEpochForCausalLM": "stablelm",
165    "MPTForCausalLM": "mpt",
166    "MT5ForConditionalGeneration": "t5",
167    "MaincoderForCausalLM": "maincoder",
168    "Mamba2ForCausalLM": "mamba",
169    "MambaForCausalLM": "mamba",
170    "MambaLMHeadModel": "mamba",
171    "MapleForCausalLM": "maple",
172    "MellumForCausalLM": "mellum",
173    "MiMoV2FlashForCausalLM": "mimo",
174    "MiMoV2ForCausalLM": "mimo",
175    "MiniCPM3ForCausalLM": "minicpm",
176    "MiniCPMForCausalLM": "minicpm",
177    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
178    "MiniMaxText01ForCausalLM": "minimax",
179    "MiniMaxM1ForCausalLM": "minimax",
180    "MiniMaxM2ForCausalLM": "minimax",
181    "MiniMaxM3SparseForCausalLM": "minimax",
182    "MiniMaxM3SparseForConditionalGeneration": "minimax",
183    "Ministral3ForCausalLM": "mistral3",
184    "Mistral3ForConditionalGeneration": "mistral3",
185    "MistralForCausalLM": "llama",
186    "MixtralForCausalLM": "llama",
187    "ModernBertForMaskedLM": "bert",
188    "ModernBertForSequenceClassification": "bert",
189    "ModernBertModel": "bert",
190    "NanbeigeForCausalLM": "nanbeige",
191    "NemotronForCausalLM": "nemotron",
192    "NemotronHForCausalLM": "nemotron",
193    "NemotronHPuzzleForCausalLM": "nemotron",
194    "NeoBERT": "bert",
195    "NeoBERTForSequenceClassification": "bert",
196    "NeoBERTLMHead": "bert",
197    "NomicBertModel": "bert",
198    "OLMoForCausalLM": "olmo",
199    "Olmo2ForCausalLM": "olmo",
200    "Olmo3ForCausalLM": "olmo",
201    "OlmoForCausalLM": "olmo",
202    "OlmoeForCausalLM": "olmo",
203    "MuseGlimmerAssistantModel": "muse_glimmer",
204    "MuseGlimmerForConditionalGeneration": "muse_glimmer",
205    "OpenELMForCausalLM": "openelm",
206    "OrionForCausalLM": "orion",
207    "PLMForCausalLM": "plm",
208    "PLaMo2ForCausalLM": "plamo",
209    "PLaMo3ForCausalLM": "plamo",
210    "PaddleOCRVLForConditionalGeneration": "ernie",
211    "PanguEmbeddedForCausalLM": "pangu",
212    "Phi3ForCausalLM": "phi",
213    "Phi4ForCausalLMV": "phi",
214    "PhiForCausalLM": "phi",
215    "PhiMoEForCausalLM": "phi",
216    "Plamo2ForCausalLM": "plamo",
217    "Plamo3ForCausalLM": "plamo",
218    "PlamoForCausalLM": "plamo",
219    "QWenLMHeadModel": "qwen",
220    "Qwen2AudioForConditionalGeneration": "qwen",
221    "Qwen2ForCausalLM": "qwen",
222    "Qwen2Model": "qwen",
223    "Qwen2MoeForCausalLM": "qwen",
224    "Qwen2VLForConditionalGeneration": "qwenvl",
225    "Qwen2VLModel": "qwenvl",
226    "Qwen2_5OmniModel": "qwenvl",
227    "Qwen2_5_VLForConditionalGeneration": "qwenvl",
228    "Qwen3ASRForConditionalGeneration": "qwen3vl",
229    "Qwen3ForCausalLM": "qwen",
230    "Qwen3Model": "qwen",
231    "Qwen3MoeForCausalLM": "qwen",
232    "Qwen3NextForCausalLM": "qwen",
233    "Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
234    "PocketTTSModel": "pockettts",
235    "Qwen3TTSForConditionalGeneration": "qwen3tts",
236    "Qwen3VLForConditionalGeneration": "qwen3vl",
237    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",
238    "Qwen3_5ForCausalLM": "qwen",
239    "Qwen3_5ForConditionalGeneration": "qwen",
240    "Qwen3_5MoeForCausalLM": "qwen",
241    "Qwen3_5MoeForConditionalGeneration": "qwen",
242    "Qwen4ExpForCausalLM": "qwen4exp",
243    "Qwen4ExpForConditionalGeneration": "qwen4exp",
244    "RND1": "qwen",
245    "RWForCausalLM": "falcon",
246    "RWKV6Qwen2ForCausalLM": "rwkv",
247    "RWKV7ForCausalLM": "rwkv",
248    "RobertaForSequenceClassification": "bert",
249    "RobertaModel": "bert",
250    "RuGPT3XLForCausalLM": "gpt2",
251    "Rwkv6ForCausalLM": "rwkv",
252    "Rwkv7ForCausalLM": "rwkv",
253    "RwkvHybridForCausalLM": "rwkv",
254    "Sarashina2VisionForCausalLM": "sarashina2",
255    "SarvamMoEForCausalLM": "bailingmoe",
256    "SeedOssForCausalLM": "olmo",
257    "SmallThinkerForCausalLM": "smallthinker",
258    "SmolLM3ForCausalLM": "llama",
259    "Spark2_5ForCausalLM": "spark2_5",
260    "SolarOpenForCausalLM": "glm",
261    "StableLMEpochForCausalLM": "stablelm",
262    "StableLmForCausalLM": "stablelm",
263    "Starcoder2ForCausalLM": "starcoder",
264    "Step3p5ForCausalLM": "step3",
265    "StepVLForConditionalGeneration": "step3",
266    "Step3p7ForConditionalGeneration": "step3",
267    "T5EncoderModel": "t5",
268    "T5ForConditionalGeneration": "t5",
269    "T5WithLMHeadModel": "t5",
270    "TalkieForCausalLM": "talkie",
271    "UMT5ForConditionalGeneration": "t5",
272    "UMT5Model": "t5",
273    "UltravoxModel": "ultravox",
274    "UnlimitedOCRForCausalLM": "deepseek",
275    "VLlama3ForCausalLM": "llama",
276    "VoxtralForConditionalGeneration": "llama",
277    "WavTokenizerDec": "wavtokenizer",
278    "XLMRobertaForSequenceClassification": "bert",
279    "XLMRobertaModel": "bert",
280    "XverseForCausalLM": "xverse",
281    "YoutuForCausalLM": "deepseek",
282    "YoutuVLForConditionalGeneration": "deepseek",
283    "modeling_grove_moe.GroveMoeForCausalLM": "grovemoe",
284    "modeling_sarvam_moe.SarvamMoEForCausalLM": "bailingmoe",
285}
286
287
288MMPROJ_MODEL_MAP: dict[str, str] = {
289    "AudioFlamingo3ForConditionalGeneration": "ultravox",
290    "CogVLMForCausalLM": "cogvlm",
291    "DeepseekOCR2ForCausalLM": "deepseek",
292    "DeepseekOCRForCausalLM": "deepseek",
293    "DeepseekV4ForCausalLM": "deepseek",
294    "Dots3NoteForCausalLM": "dots3",
295    "Dots3NoteForConditionalGeneration": "dots3",
296    "DotsOCRForCausalLM": "dotsocr",
297    "Exaone4_5_ForConditionalGeneration": "exaone",
298    "Gemma3ForConditionalGeneration": "gemma",
299    "Gemma3nForConditionalGeneration": "gemma",
300    "Gemma4ForConditionalGeneration": "gemma",
301    "Gemma4UnifiedForConditionalGeneration": "gemma",
302    "Glm4vForConditionalGeneration": "qwen3vl",
303    "Glm4vMoeForConditionalGeneration": "qwen3vl",
304    "Glm5vForConditionalGeneration": "kimivl",
305    "GlmOcrForConditionalGeneration": "qwen3vl",
306    "GlmasrModel": "ultravox",
307    "Granite4VisionForConditionalGeneration": "granite",
308    "GraniteSpeechForConditionalGeneration": "granite",
309    "GraniteSpeechPlusForConditionalGeneration": "granite",
310    "HunYuanVLForConditionalGeneration": "hunyuan",
311    "Idefics3ForConditionalGeneration": "smolvlm",
312    "InternVisionModel": "internvl",
313    "JanusForConditionalGeneration": "januspro",
314    "KimiK25ForConditionalGeneration": "kimivl",
315    "KimiVLForConditionalGeneration": "kimivl",
316    "Lfm2AudioForConditionalGeneration": "lfm2",
317    "Lfm2VlForConditionalGeneration": "lfm2",
318    "LightOnOCRForConditionalGeneration": "lighton_ocr",
319    "Llama4ForConditionalGeneration": "llama4",
320    "LlavaForConditionalGeneration": "llava",
321    "MERaLiON2ForConditionalGeneration": "ultravox",
322    "MiMoV2ForCausalLM": "mimo",
323    "MiniMaxM3SparseForConditionalGeneration": "minimax",
324    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
325    "Mistral3ForConditionalGeneration": "llava",
326    "NemotronH_Nano_VL_V2": "nemotron",
327    "MuseGlimmerForConditionalGeneration": "muse_glimmer",
328    "PaddleOCRVisionModel": "ernie",
329    "Phi4ForCausalLMV": "phi",
330    "Qwen2AudioForConditionalGeneration": "ultravox",
331    "Qwen2VLForConditionalGeneration": "qwenvl",
332    "Qwen2VLModel": "qwenvl",
333    "Qwen2_5OmniModel": "qwenvl",
334    "Qwen2_5_VLForConditionalGeneration": "qwenvl",
335    "Qwen3ASRForConditionalGeneration": "qwen3vl",
336    "Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
337    "PocketTTSModel": "pockettts",
338    "Qwen3TTSForConditionalGeneration": "qwen3tts",
339    "Qwen3VLForConditionalGeneration": "qwen3vl",
340    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",
341    "Qwen3_5ForConditionalGeneration": "qwen3vl",
342    "Qwen3_5MoeForConditionalGeneration": "qwen3vl",
343    "Qwen4ExpForConditionalGeneration": "qwen4exp",
344    "RADIOModel": "nemotron",
345    "Sarashina2VisionForCausalLM": "sarashina2",
346    "SmolVLMForConditionalGeneration": "smolvlm",
347    "StepVLForConditionalGeneration": "step3",
348    "Step3p7ForConditionalGeneration": "step3",
349    "UltravoxModel": "ultravox",
350    "UnlimitedOCRForCausalLM": "deepseek",
351    "VoxtralForConditionalGeneration": "ultravox",
352    "YoutuVLForConditionalGeneration": "youtuvl",
353}
354
355
356_TEXT_MODEL_MODULES = sorted(set(TEXT_MODEL_MAP.values()))
357_MMPROJ_MODEL_MODULES = sorted(set(MMPROJ_MODEL_MAP.values()))
358
359
360_loaded_text_modules: set[str] = set()
361_loaded_mmproj_modules: set[str] = set()
362
363
364def load_all_models() -> None:
365    """Import all model modules to trigger @ModelBase.register() decorators."""
366    if len(_loaded_text_modules) != len(_TEXT_MODEL_MODULES):
367        for module_name in _TEXT_MODEL_MODULES:
368            if module_name not in _loaded_text_modules:
369                try:
370                    __import__(f"conversion.{module_name}")
371                    _loaded_text_modules.add(module_name)
372                except Exception as e:
373                    logger.warning(f"Failed to load model module {module_name}: {e}")
374
375    if len(_loaded_mmproj_modules) != len(_MMPROJ_MODEL_MODULES):
376        for module_name in _MMPROJ_MODEL_MODULES:
377            if module_name not in _loaded_mmproj_modules:
378                try:
379                    __import__(f"conversion.{module_name}")
380                    _loaded_mmproj_modules.add(module_name)
381                except Exception as e:
382                    logger.warning(f"Failed to load model module {module_name}: {e}")
383
384
385def get_model_class(name: str, mmproj: bool = False) -> Type[ModelBase]:
386    """Dynamically import and return a model class by its HuggingFace architecture name."""
387    relevant_map = MMPROJ_MODEL_MAP if mmproj else TEXT_MODEL_MAP
388    if name not in relevant_map:
389        raise NotImplementedError(f"Architecture {name!r} not supported!")
390    module_name = relevant_map[name]
391    __import__(f"conversion.{module_name}")
392    model_type = ModelType.MMPROJ if mmproj else ModelType.TEXT
393    return ModelBase._model_classes[model_type][name]
394
395
396def print_registered_models() -> None:
397    load_all_models()
398    logger.error("TEXT models:")
399    for name in sorted(TEXT_MODEL_MAP.keys()):
400        logger.error(f"  - {name}")
401    logger.error("MMPROJ models:")
402    for name in sorted(MMPROJ_MODEL_MAP.keys()):
403        logger.error(f"  - {name}")