54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
from typing import Any
|
|
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
from transformers.models.siglip import SiglipVisionConfig
|
|
|
|
from sglang.srt.configs.jet_nemotron import JetNemotronConfig
|
|
from sglang.srt.configs.mamba_utils import Mamba2CacheParams
|
|
|
|
|
|
class JetVLMConfig(PretrainedConfig):
|
|
model_type = "jet_vlm"
|
|
sub_configs = {
|
|
"text_config": JetNemotronConfig,
|
|
"vision_config": SiglipVisionConfig,
|
|
}
|
|
_auto_class = "AutoConfig"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
text_config: dict[str, Any] | None = None,
|
|
vision_config: dict[str, Any] | None = None,
|
|
image_token_id: int | None = None,
|
|
video_token_id: int | None = None,
|
|
**kwargs,
|
|
):
|
|
self.text_config = (
|
|
JetNemotronConfig(**text_config)
|
|
if text_config is not None
|
|
else JetNemotronConfig()
|
|
)
|
|
self.vision_config = (
|
|
SiglipVisionConfig(**vision_config)
|
|
if vision_config is not None
|
|
else SiglipVisionConfig()
|
|
)
|
|
|
|
self.image_token_id = image_token_id if image_token_id is not None else -1
|
|
self.video_token_id = video_token_id if video_token_id is not None else -1
|
|
|
|
super().__init__(**kwargs)
|
|
|
|
@property
|
|
def full_attention_layer_ids(self) -> list[int]:
|
|
return self.text_config.full_attention_layer_ids
|
|
|
|
@property
|
|
def linear_layer_ids(self) -> list[int]:
|
|
return self.text_config.linear_layer_ids
|
|
|
|
@property
|
|
def mamba2_cache_params(self) -> Mamba2CacheParams:
|
|
return self.text_config.mamba2_cache_params
|