Brunobkr/llama.cpp_AlgMor24_github

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