File size: 1,811 Bytes
a596b0a | 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 | """PRISM configuration.
Defaults reproduce the released checkpoint (SigLIP2-so400m / Qwen3-Embedding-0.6B).
"""
from __future__ import annotations
from transformers import PretrainedConfig
class PRISMConfig(PretrainedConfig):
model_type = "prism"
def __init__(
self,
# ---- frozen backbones (loaded from the Hub by name) ----
vision_backbone_name: str = "google/siglip2-so400m-patch14-384",
text_backbone_name: str = "Qwen/Qwen3-Embedding-0.6B",
# ---- architecture ----
d_z: int = 512,
qformer_depth: int = 4,
temporal_depth: int = 12,
predictor_depth: int = 4,
num_heads: int = 8,
mlp_ratio: float = 4.0,
max_frames: int = 128,
logit_scale_init: float = 2.6592,
# ---- training objective ----
lambda_decomp: float = 1.0,
lambda_temp: float = 0.5,
infonce_all_gather: bool = True,
sliding_shift_aug: bool = True,
# ---- EMA target encoder ----
use_ema: bool = True,
ema_decay: float = 0.998,
**kwargs,
):
super().__init__(**kwargs)
self.vision_backbone_name = vision_backbone_name
self.text_backbone_name = text_backbone_name
self.d_z = d_z
self.qformer_depth = qformer_depth
self.temporal_depth = temporal_depth
self.predictor_depth = predictor_depth
self.num_heads = num_heads
self.mlp_ratio = mlp_ratio
self.max_frames = max_frames
self.logit_scale_init = logit_scale_init
self.lambda_decomp = lambda_decomp
self.lambda_temp = lambda_temp
self.infonce_all_gather = infonce_all_gather
self.sliding_shift_aug = sliding_shift_aug
self.use_ema = use_ema
self.ema_decay = ema_decay
|