File size: 1,918 Bytes
0ab2566 | 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 54 55 56 57 58 59 60 61 62 63 | ## TODO: 增加对于TBStarsParscale的支持
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=""
)
# sanity check
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}")
|