"""Configuration objects for the Cola2 Full-DiT model.""" from __future__ import annotations from typing import Any, Optional from transformers import GenerationConfig, PretrainedConfig _AUTO_MAP = { "AutoConfig": "configuration_cola2.Cola2Config", "AutoModel": "modeling_cola2.Cola2ForConditionalGeneration", "AutoModelForCausalLM": "modeling_cola2.Cola2ForConditionalGeneration", } class Cola2Config(PretrainedConfig): """Architecture configuration for span-causal Cola2. Defaults describe the released 4B DiT + Qwen3-4B-shaped decoder. Tests and CPU examples should use :meth:`tiny`, which executes the identical graph at a small width and depth. """ model_type = "cola2" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size: int = 151_936, encoder_embedding_size: int = 2_560, encoder_hidden_size: int = 1_536, encoder_intermediate_size: int = 6_144, encoder_num_attention_heads: int = 12, encoder_num_hidden_layers: int = 0, pool_num_hidden_layers: int = 2, latent_encoder_num_hidden_layers: int = 2, latent_decoder_num_hidden_layers: int = 0, latent_dim: int = 64, use_variational_encoder: bool = True, boundary_score_alpha: float = 10.0, boundary_score_threshold: float = 0.5, target_compression_rate: float = 0.5, min_span_length: int = 1, max_span_length: int = 8, latent_block_causal: bool = False, latent_block_size: int = 8, dit_hidden_size: int = 2_560, dit_intermediate_size: int = 10_240, dit_num_hidden_layers: int = 34, dit_num_attention_heads: int = 20, dit_head_dim: int = 128, dit_rope_dim: int = 96, dit_block_size: int = 8, decoder_hidden_size: int = 2_560, decoder_intermediate_size: int = 9_728, decoder_num_hidden_layers: int = 36, decoder_num_attention_heads: int = 32, decoder_num_key_value_heads: int = 8, decoder_head_dim: int = 128, condition_num_attention_heads: int = 20, decoder_condition_num_hidden_layers: int = 2, decoder_lm_head_bias: bool = True, duration_hidden_size: Optional[int] = None, local_position_dim: int = 128, hidden_act: str = "silu", hidden_dropout: float = 0.0, attention_dropout: float = 0.0, layer_norm_eps: float = 1e-6, dit_layer_norm_eps: float = 1e-5, encoder_rope_theta: float = 500_000.0, condition_rope_theta: float = 500_000.0, decoder_rope_theta: float = 1_000_000.0, dit_rope_theta: float = 10_000.0, rope_theta: float = 1_000_000.0, initializer_range: float = 0.02, max_position_embeddings: int = 32_768, scaling_factor: float = 1.0, shifting_factor: float = 0.0, diffusion_time: float = 1_000.0, pad_token_id: int = 151_936, mask_token_id: int = 151_936, bos_token_id: Optional[int] = 151_643, eos_token_id: Optional[int] = 151_643, tie_word_embeddings: bool = False, auto_map: Optional[dict[str, str]] = None, architectures: Optional[list[str]] = None, **kwargs: Any, ) -> None: self.vocab_size = int(vocab_size) self.encoder_embedding_size = int(encoder_embedding_size) self.encoder_hidden_size = int(encoder_hidden_size) self.encoder_intermediate_size = int(encoder_intermediate_size) self.encoder_num_attention_heads = int(encoder_num_attention_heads) self.encoder_num_hidden_layers = int(encoder_num_hidden_layers) self.pool_num_hidden_layers = int(pool_num_hidden_layers) self.latent_encoder_num_hidden_layers = int( latent_encoder_num_hidden_layers ) self.latent_decoder_num_hidden_layers = int( latent_decoder_num_hidden_layers ) self.latent_dim = int(latent_dim) self.use_variational_encoder = bool(use_variational_encoder) self.boundary_score_alpha = float(boundary_score_alpha) self.boundary_score_threshold = float(boundary_score_threshold) self.target_compression_rate = float(target_compression_rate) self.min_span_length = int(min_span_length) self.max_span_length = int(max_span_length) self.latent_block_causal = bool(latent_block_causal) self.latent_block_size = int(latent_block_size) self.dit_hidden_size = int(dit_hidden_size) self.dit_intermediate_size = int(dit_intermediate_size) self.dit_num_hidden_layers = int(dit_num_hidden_layers) self.dit_num_attention_heads = int(dit_num_attention_heads) self.dit_head_dim = int(dit_head_dim) self.dit_rope_dim = int(dit_rope_dim) self.dit_block_size = int(dit_block_size) self.decoder_hidden_size = int(decoder_hidden_size) self.decoder_intermediate_size = int(decoder_intermediate_size) self.decoder_num_hidden_layers = int(decoder_num_hidden_layers) self.decoder_num_attention_heads = int(decoder_num_attention_heads) self.decoder_num_key_value_heads = int(decoder_num_key_value_heads) self.decoder_head_dim = int(decoder_head_dim) self.condition_num_attention_heads = int( condition_num_attention_heads ) self.decoder_condition_num_hidden_layers = int( decoder_condition_num_hidden_layers ) self.decoder_lm_head_bias = bool(decoder_lm_head_bias) self.duration_hidden_size = int( duration_hidden_size or max(self.decoder_hidden_size // 2, 64) ) self.local_position_dim = int(local_position_dim) self.hidden_act = str(hidden_act) self.hidden_dropout = float(hidden_dropout) self.attention_dropout = float(attention_dropout) self.layer_norm_eps = float(layer_norm_eps) self.dit_layer_norm_eps = float(dit_layer_norm_eps) self.encoder_rope_theta = float(encoder_rope_theta) self.condition_rope_theta = float(condition_rope_theta) self.decoder_rope_theta = float(decoder_rope_theta) self.dit_rope_theta = float(dit_rope_theta) self.rope_theta = float(rope_theta) self.initializer_range = float(initializer_range) self.max_position_embeddings = int(max_position_embeddings) self.scaling_factor = float(scaling_factor) self.shifting_factor = float(shifting_factor) self.diffusion_time = float(diffusion_time) self.mask_token_id = int(mask_token_id) self._validate_architecture() # These are serialized by PretrainedConfig and are fixed by this model. # Remove them from **kwargs to avoid passing the same keyword twice on a # config.json roundtrip. kwargs.pop("is_decoder", None) kwargs.pop("is_encoder_decoder", None) super().__init__( pad_token_id=int(pad_token_id), mask_token_id=int(mask_token_id), bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=bool(tie_word_embeddings), is_decoder=True, is_encoder_decoder=False, auto_map=dict(_AUTO_MAP if auto_map is None else auto_map), architectures=( ["Cola2ForConditionalGeneration"] if architectures is None else architectures ), **kwargs, ) def _validate_architecture(self) -> None: positive = { "vocab_size": self.vocab_size, "encoder_embedding_size": self.encoder_embedding_size, "encoder_hidden_size": self.encoder_hidden_size, "encoder_intermediate_size": self.encoder_intermediate_size, "encoder_num_attention_heads": self.encoder_num_attention_heads, "latent_dim": self.latent_dim, "latent_block_size": self.latent_block_size, "dit_hidden_size": self.dit_hidden_size, "dit_intermediate_size": self.dit_intermediate_size, "dit_num_attention_heads": self.dit_num_attention_heads, "dit_head_dim": self.dit_head_dim, "dit_block_size": self.dit_block_size, "decoder_hidden_size": self.decoder_hidden_size, "decoder_intermediate_size": self.decoder_intermediate_size, "decoder_num_attention_heads": self.decoder_num_attention_heads, "decoder_num_key_value_heads": self.decoder_num_key_value_heads, "decoder_head_dim": self.decoder_head_dim, "condition_num_attention_heads": self.condition_num_attention_heads, "max_span_length": self.max_span_length, } for name, value in positive.items(): if value <= 0: raise ValueError(f"{name} must be positive, got {value}.") for name in ( "encoder_num_hidden_layers", "pool_num_hidden_layers", "latent_encoder_num_hidden_layers", "latent_decoder_num_hidden_layers", "dit_num_hidden_layers", "decoder_num_hidden_layers", "decoder_condition_num_hidden_layers", ): if getattr(self, name) < 0: raise ValueError(f"{name} cannot be negative.") if self.encoder_hidden_size % self.encoder_num_attention_heads: raise ValueError( "encoder_hidden_size must be divisible by " "encoder_num_attention_heads." ) if ( self.dit_num_attention_heads * self.dit_head_dim != self.dit_hidden_size ): raise ValueError( "dit_num_attention_heads * dit_head_dim must equal " "dit_hidden_size." ) if ( self.decoder_num_attention_heads * self.decoder_head_dim <= 0 ): raise ValueError("Invalid decoder attention geometry.") if ( self.decoder_num_attention_heads % self.decoder_num_key_value_heads ): raise ValueError( "decoder_num_attention_heads must be divisible by " "decoder_num_key_value_heads." ) if self.decoder_hidden_size % self.condition_num_attention_heads: raise ValueError( "decoder_hidden_size must be divisible by " "condition_num_attention_heads." ) if not 0 < self.dit_rope_dim <= self.dit_head_dim: raise ValueError("dit_rope_dim must be in (0, dit_head_dim].") if self.dit_rope_dim % 2: raise ValueError("dit_rope_dim must be even.") if self.decoder_head_dim % 2: raise ValueError("decoder_head_dim must be even.") if not 1 <= self.min_span_length <= self.max_span_length: raise ValueError( "Span lengths must satisfy 1 <= min_span_length <= " "max_span_length." ) if not 0 < self.target_compression_rate <= 1: raise ValueError("target_compression_rate must be in (0, 1].") if self.scaling_factor == 0: raise ValueError("scaling_factor cannot be zero.") if self.diffusion_time <= 0: raise ValueError("diffusion_time must be positive.") @classmethod def tiny(cls, **overrides: Any) -> "Cola2Config": """Return a CPU-friendly config that uses the production code path.""" values: dict[str, Any] = { "vocab_size": 67, "encoder_embedding_size": 32, "encoder_hidden_size": 32, "encoder_intermediate_size": 64, "encoder_num_attention_heads": 4, "encoder_num_hidden_layers": 0, "pool_num_hidden_layers": 1, "latent_encoder_num_hidden_layers": 1, "latent_decoder_num_hidden_layers": 0, "latent_dim": 8, "min_span_length": 1, "max_span_length": 3, "latent_block_size": 4, "dit_hidden_size": 32, "dit_intermediate_size": 64, "dit_num_hidden_layers": 2, "dit_num_attention_heads": 4, "dit_head_dim": 8, "dit_rope_dim": 8, "dit_block_size": 4, "decoder_hidden_size": 32, "decoder_intermediate_size": 64, "decoder_num_hidden_layers": 2, "decoder_num_attention_heads": 4, "decoder_num_key_value_heads": 2, "decoder_head_dim": 8, "condition_num_attention_heads": 4, "decoder_condition_num_hidden_layers": 1, "duration_hidden_size": 32, "local_position_dim": 8, "max_position_embeddings": 256, "pad_token_id": 0, "mask_token_id": 0, "bos_token_id": 1, "eos_token_id": 2, } values.update(overrides) return cls(**values) class Cola2GenerationConfig(GenerationConfig): """Generation controls for the Full-DiT + span-causal inference path.""" model_type = "cola2" def __init__( self, dit_num_steps: int = 8, dit_timestep_schedule: str = "low_noise", dit_timestep_power: Optional[float] = None, dit_cfg_disable_empty_prefix: bool = False, dit_guidance_scale: float = 10.0, decoder_guidance_scale: float = -0.9, span_causal_chunk_size: int = 1, dynamic_duration: bool = True, latent_feedback: bool = True, **kwargs: Any, ) -> None: # Accepted aliases mirror names used by the original evaluation harness. dit_num_steps = int(kwargs.pop("dit_steps", dit_num_steps)) dit_guidance_scale = float( kwargs.pop("dit_cfg_scale", dit_guidance_scale) ) decoder_guidance_scale = float( kwargs.pop("decoder_cfg_scale", decoder_guidance_scale) ) span_causal_chunk_size = int( kwargs.pop( "span_causal_chunk_spans", kwargs.pop("chunk_spans", span_causal_chunk_size), ) ) latent_feedback = bool( kwargs.pop("latent_feedback_reencode", latent_feedback) ) kwargs.setdefault("max_new_tokens", 32) kwargs.setdefault("do_sample", False) kwargs.setdefault("temperature", 1.0) kwargs.setdefault("top_k", 50) kwargs.setdefault("top_p", 1.0) super().__init__(**kwargs) self.dit_num_steps = dit_num_steps self.dit_timestep_schedule = str(dit_timestep_schedule) self.dit_timestep_power = ( None if dit_timestep_power is None else float(dit_timestep_power) ) self.dit_cfg_disable_empty_prefix = bool( dit_cfg_disable_empty_prefix ) self.dit_guidance_scale = dit_guidance_scale self.decoder_guidance_scale = decoder_guidance_scale self.span_causal_chunk_size = span_causal_chunk_size self.dynamic_duration = bool(dynamic_duration) self.latent_feedback = bool(latent_feedback) self.validate() @property def dit_steps(self) -> int: return self.dit_num_steps @dit_steps.setter def dit_steps(self, value: int) -> None: self.dit_num_steps = int(value) @property def dit_cfg_scale(self) -> float: return self.dit_guidance_scale @dit_cfg_scale.setter def dit_cfg_scale(self, value: float) -> None: self.dit_guidance_scale = float(value) @property def decoder_cfg_scale(self) -> float: return self.decoder_guidance_scale @decoder_cfg_scale.setter def decoder_cfg_scale(self, value: float) -> None: self.decoder_guidance_scale = float(value) @property def chunk_spans(self) -> int: return self.span_causal_chunk_size @chunk_spans.setter def chunk_spans(self, value: int) -> None: self.span_causal_chunk_size = int(value) @property def span_causal_chunk_spans(self) -> int: return self.span_causal_chunk_size @span_causal_chunk_spans.setter def span_causal_chunk_spans(self, value: int) -> None: self.span_causal_chunk_size = int(value) @property def latent_feedback_reencode(self) -> bool: return self.latent_feedback @latent_feedback_reencode.setter def latent_feedback_reencode(self, value: bool) -> None: self.latent_feedback = bool(value) def validate(self, is_init: bool = False) -> None: super().validate(is_init=is_init) if getattr(self, "dit_num_steps", 1) < 1: raise ValueError("dit_num_steps must be at least 1.") if getattr(self, "dit_timestep_schedule", "uniform") not in { "uniform", "high_noise", "low_noise", "power", }: raise ValueError( "dit_timestep_schedule must be uniform, high_noise, " "low_noise, or power." ) if ( getattr(self, "dit_timestep_schedule", "uniform") == "power" and ( getattr(self, "dit_timestep_power", None) is None or float(self.dit_timestep_power) <= 0 ) ): raise ValueError( "dit_timestep_power must be positive for power schedule." ) if getattr(self, "span_causal_chunk_size", 1) < 1: raise ValueError("span_causal_chunk_size must be at least 1.")