|
|
from dataclasses import dataclass, field
|
|
|
from typing import Dict, List, Optional
|
|
|
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
|
class TextConfig:
|
|
|
dim: int = 2048
|
|
|
ff_dim: int = 8192
|
|
|
n_layers: int = 24
|
|
|
vocab_size: int = 51200
|
|
|
max_context: int = 2048
|
|
|
n_heads: int = 32
|
|
|
n_kv_heads: int = 32
|
|
|
prefix_attn: int = 730
|
|
|
group_size: Optional[int] = None
|
|
|
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
|
class VisionConfig:
|
|
|
enc_dim: int = 1152
|
|
|
enc_patch_size: int = 14
|
|
|
enc_n_layers: int = 27
|
|
|
enc_ff_dim: int = 4304
|
|
|
enc_n_heads: int = 16
|
|
|
proj_out_dim: int = 2048
|
|
|
crop_size: int = 378
|
|
|
in_channels: int = 3
|
|
|
max_crops: int = 12
|
|
|
overlap_margin: int = 4
|
|
|
proj_inner_dim: int = 8192
|
|
|
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
|
class RegionConfig:
|
|
|
dim: int = 2048
|
|
|
coord_feat_dim: int = 256
|
|
|
coord_out_dim: int = 1024
|
|
|
size_feat_dim: int = 512
|
|
|
size_out_dim: int = 2048
|
|
|
inner_dim: int = 8192
|
|
|
group_size: Optional[int] = None
|
|
|
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
|
class TokenizerConfig:
|
|
|
bos_id: int = 0
|
|
|
eos_id: int = 0
|
|
|
answer_id: int = 3
|
|
|
thinking_id: int = 4
|
|
|
coord_id: int = 5
|
|
|
size_id: int = 6
|
|
|
start_ground_points_id: int = 7
|
|
|
end_ground_id: int = 9
|
|
|
templates: Dict[str, Optional[Dict[str, List[int]]]] = field(
|
|
|
default_factory=lambda: {
|
|
|
"caption": {
|
|
|
"short": [1, 32708, 2, 12492, 3],
|
|
|
"normal": [1, 32708, 2, 6382, 3],
|
|
|
"long": [1, 32708, 2, 4059, 3],
|
|
|
},
|
|
|
"query": {"prefix": [1, 15381, 2], "suffix": [3]},
|
|
|
"detect": {"prefix": [1, 7235, 476, 2], "suffix": [3]},
|
|
|
"point": {"prefix": [1, 2581, 2], "suffix": [3]},
|
|
|
}
|
|
|
)
|
|
|
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
|
class MoondreamConfig:
|
|
|
text: TextConfig = TextConfig()
|
|
|
vision: VisionConfig = VisionConfig()
|
|
|
region: RegionConfig = RegionConfig()
|
|
|
tokenizer: TokenizerConfig = TokenizerConfig()
|
|
|
|
|
|
@classmethod
|
|
|
def from_dict(cls, config_dict: dict):
|
|
|
text_config = TextConfig(**config_dict.get("text", {}))
|
|
|
vision_config = VisionConfig(**config_dict.get("vision", {}))
|
|
|
region_config = RegionConfig(**config_dict.get("region", {}))
|
|
|
tokenizer_config = TokenizerConfig(**config_dict.get("tokenizer", {}))
|
|
|
return cls(
|
|
|
text=text_config,
|
|
|
vision=vision_config,
|
|
|
region=region_config,
|
|
|
tokenizer=tokenizer_config,
|
|
|
)
|
|
|
|
|
|
def to_dict(self):
|
|
|
return {
|
|
|
"text": self.text.__dict__,
|
|
|
"vision": self.vision.__dict__,
|
|
|
"region": self.region.__dict__,
|
|
|
"tokenizer": self.tokenizer.__dict__,
|
|
|
}
|
|
|
|
