| 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}" |
|
|
