Basic script for OLMoE inference.

Add: basic script for OLMoE inference and saving expert logits.

Add: OLMoE vLLM module with logits saving.

Add: vLLM model registering.

Fix: update DataLogger module for multiprocessing.

models/register.py

@@ -0,0 +1,31 @@
+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"
+    # )