Felipe97/llama-cpp-compiled

1from __future__ import annotations2 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 Type9 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    "HrmTextForCausalLM": "hrm_text",127    "HYV3ForCausalLM": "hunyuan",128    "HYV4ForCausalLM": "hy_v4",129    "IQuestCoderForCausalLM": "llama",130    "InternLM2ForCausalLM": "internlm",131    "InternLM3ForCausalLM": "internlm",132    "JAISLMHeadModel": "jais",133    "Jais2ForCausalLM": "jais",134    "JambaForCausalLM": "jamba",135    "JanusForConditionalGeneration": "januspro",136    "JinaBertForMaskedLM": "bert",137    "JinaBertModel": "bert",138    "JinaEmbeddingsV5Model": "bert",139    "KORMoForCausalLM": "qwen",140    "KimiK25ForConditionalGeneration": "deepseek",141    "KimiK3ForConditionalGeneration": "kimi_k3",142    "KimiLinearForCausalLM": "kimi_linear",143    "KimiLinearModel": "kimi_linear",144    "KimiVLForConditionalGeneration": "deepseek",145    "LFM2ForCausalLM": "lfm2",146    "LLaDAMoEModel": "llada",147    "LLaDAMoEModelLM": "llada",148    "LLaDAModelLM": "llada",149    "LLaMAForCausalLM": "llama",150    "Lfm25AudioTokenizer": "lfm2",151    "Lfm2BidirectionalModel": "lfm2",152    "Lfm2ForCausalLM": "lfm2",153    "Lfm2Model": "lfm2",154    "Lfm2MoeForCausalLM": "lfm2",155    "Llama4ForCausalLM": "llama",156    "Llama4ForConditionalGeneration": "llama",157    "LlamaBidirectionalModel": "llama",158    "LlamaForCausalLM": "llama",159    "LlamaModel": "llama",160    "Eagle3DraftModel": "llama",161    "Eagle3Speculator": "llama",162    "Eagle3LlamaForCausalLM": "llama",163    "LlamaForCausalLMEagle3": "llama",164    "LlavaForConditionalGeneration": "llama",165    "LlavaStableLMEpochForCausalLM": "stablelm",166    "MPTForCausalLM": "mpt",167    "MT5ForConditionalGeneration": "t5",168    "MaincoderForCausalLM": "maincoder",169    "Mamba2ForCausalLM": "mamba",170    "MambaForCausalLM": "mamba",171    "MambaLMHeadModel": "mamba",172    "MapleForCausalLM": "maple",173    "MellumForCausalLM": "mellum",174    "MiMoV2FlashForCausalLM": "mimo",175    "MiMoV2ForCausalLM": "mimo",176    "MiniCPM3ForCausalLM": "minicpm",177    "MiniCPMForCausalLM": "minicpm",178    "MiniCPMV4_6ForConditionalGeneration": "minicpm",179    "MiniMaxText01ForCausalLM": "minimax",180    "MiniMaxM1ForCausalLM": "minimax",181    "MiniMaxM2ForCausalLM": "minimax",182    "MiniMaxM3SparseForCausalLM": "minimax",183    "MiniMaxM3SparseForConditionalGeneration": "minimax",184    "Ministral3ForCausalLM": "mistral3",185    "Mistral3ForConditionalGeneration": "mistral3",186    "MistralForCausalLM": "llama",187    "MixtralForCausalLM": "llama",188    "ModernBertForMaskedLM": "bert",189    "ModernBertForSequenceClassification": "bert",190    "ModernBertModel": "bert",191    "NanbeigeForCausalLM": "nanbeige",192    "NemotronForCausalLM": "nemotron",193    "NemotronHForCausalLM": "nemotron",194    "NemotronHPuzzleForCausalLM": "nemotron",195    "NeoBERT": "bert",196    "NeoBERTForSequenceClassification": "bert",197    "NeoBERTLMHead": "bert",198    "NomicBertModel": "bert",199    "OLMoForCausalLM": "olmo",200    "Olmo2ForCausalLM": "olmo",201    "Olmo3ForCausalLM": "olmo",202    "OlmoForCausalLM": "olmo",203    "OlmoeForCausalLM": "olmo",204    "MuseGlimmerAssistantModel": "muse_glimmer",205    "MuseGlimmerForConditionalGeneration": "muse_glimmer",206    "OpenELMForCausalLM": "openelm",207    "OrionForCausalLM": "orion",208    "PLMForCausalLM": "plm",209    "PLaMo2ForCausalLM": "plamo",210    "PLaMo3ForCausalLM": "plamo",211    "PaddleOCRVLForConditionalGeneration": "ernie",212    "PanguEmbeddedForCausalLM": "pangu",213    "Phi3ForCausalLM": "phi",214    "Phi4ForCausalLMV": "phi",215    "PhiForCausalLM": "phi",216    "PhiMoEForCausalLM": "phi",217    "Plamo2ForCausalLM": "plamo",218    "Plamo3ForCausalLM": "plamo",219    "PlamoForCausalLM": "plamo",220    "QWenLMHeadModel": "qwen",221    "Qwen2AudioForConditionalGeneration": "qwen",222    "Qwen2ForCausalLM": "qwen",223    "Qwen2Model": "qwen",224    "Qwen2MoeForCausalLM": "qwen",225    "Qwen2VLForConditionalGeneration": "qwenvl",226    "Qwen2VLModel": "qwenvl",227    "Qwen2_5OmniModel": "qwenvl",228    "Qwen2_5_VLForConditionalGeneration": "qwenvl",229    "Qwen3ASRForConditionalGeneration": "qwen3vl",230    "Qwen3ForCausalLM": "qwen",231    "Qwen3Model": "qwen",232    "Qwen3MoeForCausalLM": "qwen",233    "Qwen3NextForCausalLM": "qwen",234    "Qwen3OmniMoeForConditionalGeneration": "qwen3vl",235    "PocketTTSModel": "pockettts",236    "Qwen3TTSForConditionalGeneration": "qwen3tts",237    "Qwen3VLForConditionalGeneration": "qwen3vl",238    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",239    "Qwen3_5ForCausalLM": "qwen",240    "Qwen3_5ForConditionalGeneration": "qwen",241    "Qwen3_5MoeForCausalLM": "qwen",242    "Qwen3_5MoeForConditionalGeneration": "qwen",243    "Qwen4ExpForCausalLM": "qwen4exp",244    "Qwen4ExpForConditionalGeneration": "qwen4exp",245    "RND1": "qwen",246    "RWForCausalLM": "falcon",247    "RWKV6Qwen2ForCausalLM": "rwkv",248    "RWKV7ForCausalLM": "rwkv",249    "RobertaForSequenceClassification": "bert",250    "RobertaModel": "bert",251    "RuGPT3XLForCausalLM": "gpt2",252    "Rwkv6ForCausalLM": "rwkv",253    "Rwkv7ForCausalLM": "rwkv",254    "RwkvHybridForCausalLM": "rwkv",255    "Sarashina2VisionForCausalLM": "sarashina2",256    "SarvamMoEForCausalLM": "bailingmoe",257    "SeedOssForCausalLM": "olmo",258    "SmallThinkerForCausalLM": "smallthinker",259    "SmolLM3ForCausalLM": "llama",260    "Spark2_5ForCausalLM": "spark2_5",261    "SolarOpenForCausalLM": "glm",262    "StableLMEpochForCausalLM": "stablelm",263    "StableLmForCausalLM": "stablelm",264    "Starcoder2ForCausalLM": "starcoder",265    "Step3p5ForCausalLM": "step3",266    "StepVLForConditionalGeneration": "step3",267    "Step3p7ForConditionalGeneration": "step3",268    "T5EncoderModel": "t5",269    "T5ForConditionalGeneration": "t5",270    "T5WithLMHeadModel": "t5",271    "TalkieForCausalLM": "talkie",272    "UMT5ForConditionalGeneration": "t5",273    "UMT5Model": "t5",274    "UltravoxModel": "ultravox",275    "UnlimitedOCRForCausalLM": "deepseek",276    "VLlama3ForCausalLM": "llama",277    "VoxtralForConditionalGeneration": "llama",278    "WavTokenizerDec": "wavtokenizer",279    "XLMRobertaForSequenceClassification": "bert",280    "XLMRobertaModel": "bert",281    "XverseForCausalLM": "xverse",282    "YoutuForCausalLM": "deepseek",283    "YoutuVLForConditionalGeneration": "deepseek",284    "modeling_grove_moe.GroveMoeForCausalLM": "grovemoe",285    "modeling_sarvam_moe.SarvamMoEForCausalLM": "bailingmoe",286}287 288 289MMPROJ_MODEL_MAP: dict[str, str] = {290    "AudioFlamingo3ForConditionalGeneration": "ultravox",291    "CogVLMForCausalLM": "cogvlm",292    "DeepseekOCR2ForCausalLM": "deepseek",293    "DeepseekOCRForCausalLM": "deepseek",294    "DeepseekV4ForCausalLM": "deepseek",295    "Dots3NoteForCausalLM": "dots3",296    "Dots3NoteForConditionalGeneration": "dots3",297    "DotsOCRForCausalLM": "dotsocr",298    "Exaone4_5_ForConditionalGeneration": "exaone",299    "Gemma3ForConditionalGeneration": "gemma",300    "Gemma3nForConditionalGeneration": "gemma",301    "Gemma4ForConditionalGeneration": "gemma",302    "Gemma4UnifiedForConditionalGeneration": "gemma",303    "Glm4vForConditionalGeneration": "qwen3vl",304    "Glm4vMoeForConditionalGeneration": "qwen3vl",305    "Glm5vForConditionalGeneration": "kimivl",306    "GlmOcrForConditionalGeneration": "qwen3vl",307    "GlmasrModel": "ultravox",308    "Granite4VisionForConditionalGeneration": "granite",309    "GraniteSpeechForConditionalGeneration": "granite",310    "GraniteSpeechPlusForConditionalGeneration": "granite",311    "HunYuanVLForConditionalGeneration": "hunyuan",312    "Idefics3ForConditionalGeneration": "smolvlm",313    "InternVisionModel": "internvl",314    "JanusForConditionalGeneration": "januspro",315    "KimiK25ForConditionalGeneration": "kimivl",316    "KimiVLForConditionalGeneration": "kimivl",317    "Lfm2AudioForConditionalGeneration": "lfm2",318    "Lfm2VlForConditionalGeneration": "lfm2",319    "LightOnOCRForConditionalGeneration": "lighton_ocr",320    "Llama4ForConditionalGeneration": "llama4",321    "LlavaForConditionalGeneration": "llava",322    "MERaLiON2ForConditionalGeneration": "ultravox",323    "MiMoV2ForCausalLM": "mimo",324    "MiniMaxM3SparseForConditionalGeneration": "minimax",325    "MiniCPMV4_6ForConditionalGeneration": "minicpm",326    "Mistral3ForConditionalGeneration": "llava",327    "NemotronH_Nano_VL_V2": "nemotron",328    "MuseGlimmerForConditionalGeneration": "muse_glimmer",329    "PaddleOCRVisionModel": "ernie",330    "Phi4ForCausalLMV": "phi",331    "Qwen2AudioForConditionalGeneration": "ultravox",332    "Qwen2VLForConditionalGeneration": "qwenvl",333    "Qwen2VLModel": "qwenvl",334    "Qwen2_5OmniModel": "qwenvl",335    "Qwen2_5_VLForConditionalGeneration": "qwenvl",336    "Qwen3ASRForConditionalGeneration": "qwen3vl",337    "Qwen3OmniMoeForConditionalGeneration": "qwen3vl",338    "PocketTTSModel": "pockettts",339    "Qwen3TTSForConditionalGeneration": "qwen3tts",340    "Qwen3VLForConditionalGeneration": "qwen3vl",341    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",342    "Qwen3_5ForConditionalGeneration": "qwen3vl",343    "Qwen3_5MoeForConditionalGeneration": "qwen3vl",344    "Qwen4ExpForConditionalGeneration": "qwen4exp",345    "RADIOModel": "nemotron",346    "Sarashina2VisionForCausalLM": "sarashina2",347    "SmolVLMForConditionalGeneration": "smolvlm",348    "StepVLForConditionalGeneration": "step3",349    "Step3p7ForConditionalGeneration": "step3",350    "UltravoxModel": "ultravox",351    "UnlimitedOCRForCausalLM": "deepseek",352    "VoxtralForConditionalGeneration": "ultravox",353    "YoutuVLForConditionalGeneration": "youtuvl",354}355 356 357_TEXT_MODEL_MODULES = sorted(set(TEXT_MODEL_MAP.values()))358_MMPROJ_MODEL_MODULES = sorted(set(MMPROJ_MODEL_MAP.values()))359 360 361_loaded_text_modules: set[str] = set()362_loaded_mmproj_modules: set[str] = set()363 364 365def load_all_models() -> None:366    """Import all model modules to trigger @ModelBase.register() decorators."""367    if len(_loaded_text_modules) != len(_TEXT_MODEL_MODULES):368        for module_name in _TEXT_MODEL_MODULES:369            if module_name not in _loaded_text_modules:370                try:371                    __import__(f"conversion.{module_name}")372                    _loaded_text_modules.add(module_name)373                except Exception as e:374                    logger.warning(f"Failed to load model module {module_name}: {e}")375 376    if len(_loaded_mmproj_modules) != len(_MMPROJ_MODEL_MODULES):377        for module_name in _MMPROJ_MODEL_MODULES:378            if module_name not in _loaded_mmproj_modules:379                try:380                    __import__(f"conversion.{module_name}")381                    _loaded_mmproj_modules.add(module_name)382                except Exception as e:383                    logger.warning(f"Failed to load model module {module_name}: {e}")384 385 386def get_model_class(name: str, mmproj: bool = False) -> Type[ModelBase]:387    """Dynamically import and return a model class by its HuggingFace architecture name."""388    relevant_map = MMPROJ_MODEL_MAP if mmproj else TEXT_MODEL_MAP389    if name not in relevant_map:390        raise NotImplementedError(f"Architecture {name!r} not supported!")391    module_name = relevant_map[name]392    __import__(f"conversion.{module_name}")393    model_type = ModelType.MMPROJ if mmproj else ModelType.TEXT394    return ModelBase._model_classes[model_type][name]395 396 397def print_registered_models() -> None:398    load_all_models()399    logger.error("TEXT models:")400    for name in sorted(TEXT_MODEL_MAP.keys()):401        logger.error(f"  - {name}")402    logger.error("MMPROJ models:")403    for name in sorted(MMPROJ_MODEL_MAP.keys()):404        logger.error(f"  - {name}")405