File size: 2,647 Bytes
53b1a83 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | 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__,
}
|