from vllm import ModelRegistry


def register_vllm_logit_logging_models():
    models_mapper = {"OlmoeForCausalLM": "models.log_expert.olmoe:OlmoeForCausalLM"}
    for model_name, model_path in models_mapper.items():
        ModelRegistry.register_model(model_name, model_path)

    # ModelRegistry.register_model(
    #     "Qwen3MoeForCausalLM", "src.modeling_vllm_save.qwen3_moe:Qwen3MoeForCausalLM"
    # )
    # ModelRegistry.register_model(
    #     "MixtralForCausalLM", "src.modeling_vllm_save.mixtral:MixtralForCausalLM"
    # )
    # ModelRegistry.register_model(
    #     "OlmoeForCausalLM", "models.log_expert.olmoe:OlmoeForCausalLM"
    # )
    # ModelRegistry.register_model(
    #     "DeepseekV2ForCausalLM",
    #     "src.modeling_vllm_save.deepseek_v2:DeepseekV2ForCausalLM",
    # )
    # ModelRegistry.register_model(
    #     "Llama4ForConditionalGeneration",
    #     "src.modeling_vllm_save.mllama4:Llama4ForConditionalGeneration",
    # )
    # ModelRegistry.register_model(
    #     "GptOssForCausalLM", "src.modeling_vllm_save.gpt_oss:GptOssForCausalLM"
    # )
    # ModelRegistry.register_model(
    #     "PhiMoEForCausalLM", "src.modeling_vllm_save.phimoe:PhiMoEForCausalLM"
    # )