File size: 1,684 Bytes
01cf98b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | from typing import Dict, List
from collections import OrderedDict
from collators import COLLATORS
from loaders import LOADERS
MODULE_KEYWORDS: Dict[str, Dict[str, List]] = {
"qwen3-vl": {
"vision_encoder": ["visual.patch_embed", "visual.rotary_pos_emb", "visual.blocks"],
"vision_projector": ["visual.merger"],
"llm": ["language_model"]
},
"qwen3_5": {
"vision_encoder": ["model.visual.patch_embed", "model.visual.pos_embed", "model.visual.blocks"],
"vision_projector": ["model.visual.merger"],
"llm": ["model.language_model"]
},
}
MODEL_HF_PATH = OrderedDict()
MODEL_FAMILIES = OrderedDict()
def register_model(model_id: str, model_family_id: str, model_hf_path: str) -> None:
if model_id in MODEL_HF_PATH or model_id in MODEL_FAMILIES:
raise ValueError(f"Duplicate model_id: {model_id}")
MODEL_HF_PATH[model_id] = model_hf_path
MODEL_FAMILIES[model_id] = model_family_id
register_model(
model_id="qwen3-vl-2b",
model_family_id="qwen3-vl",
model_hf_path="Qwen/Qwen3-VL-2B-Instruct"
)
register_model(
model_id="qwen3-vl-8b",
model_family_id="qwen3-vl",
model_hf_path="Qwen/Qwen3-VL-8B-Instruct"
)
register_model(
model_id="qwen3_5-9b",
model_family_id="qwen3_5",
model_hf_path="Qwen/Qwen3.5-VL-9B"
)
for model_family_id in MODEL_FAMILIES.values():
assert model_family_id in COLLATORS, f"Collator not found for model family: {model_family_id}"
assert model_family_id in LOADERS, f"Loader not found for model family: {model_family_id}"
assert model_family_id in MODULE_KEYWORDS, f"Module keywords not found for model family: {model_family_id}"
|