| |
| from typing import Dict, List |
| from collections import OrderedDict |
|
|
| from collators import COLLATORS |
| from loaders import LOADERS |
|
|
|
|
| MODULE_KEYWORDS: Dict[str, Dict[str, List]] = { |
| "tbstars_ecom_emb_3b": { |
| "vision_encoder": ["model.vision_tower"], |
| "vision_projector": ["model.mm_projector"], |
| "llm": ["model.model", "model.lm_head"] |
| }, |
| "qwen3vl": { |
| "vision_encoder": ["model.visual"], |
| "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="tbstars_ecom_emb_3b", |
| model_family_id="tbstars_ecom_emb_3b", |
| model_hf_path="nebula_internal.TBStars-VL-Ecom-EMB/version=v1" |
| ) |
|
|
| register_model( |
| model_id="qwen3vl", |
| model_family_id="qwen3vl", |
| model_hf_path="" |
| ) |
|
|
|
|
| |
| for model_family_id in MODEL_FAMILIES.values(): |
| if model_family_id != "qwen3vl" and model_family_id not in LOADERS: |
| import warnings |
| warnings.warn(f"Loader not found for model family: {model_family_id}, it will be unavailable") |
| assert model_family_id in MODULE_KEYWORDS, f"Module keywords not found for model family: {model_family_id}" |
|
|
|
|
| if __name__ == "__main__": |
| temp = "Model ID" |
| ljust = 30 |
| print("Supported models:") |
| print(f" {temp.ljust(ljust)}: HuggingFace Path") |
| print(" ------------------------------------------------") |
| for model_id, model_hf_path in MODEL_HF_PATH.items(): |
| print(f" {model_id.ljust(ljust)}: {model_hf_path}") |
|
|
