| |
| |
| """Configuration classes for Modilify Mk1.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Any |
|
|
| from transformers.models.diffusion_gemma import ( |
| DiffusionGemmaConfig, |
| DiffusionGemmaTextConfig, |
| ) |
|
|
|
|
| class ModilifyMk1TextConfig(DiffusionGemmaTextConfig): |
| """Text configuration for the Modilify Mk1 decoder. |
| |
| This class preserves the standard DiffusionGemma text schema while giving |
| the exported model an independent, stable model type. |
| """ |
|
|
| model_type = "modilify_mk1_text" |
|
|
|
|
| class ModilifyMk1Config(DiffusionGemmaConfig): |
| """Serializable multimodal inference configuration for Modilify Mk1. |
| |
| Args: |
| text_config: DiffusionGemma text configuration or its serialized form. |
| vision_config: Gemma 4 vision configuration or its serialized form. |
| denoise_temperature: Sampling temperature used at every denoising step. |
| commit_failure_budget: Maximum cumulative failure risk for normal commits. |
| fused_entropy_weight: Multiplicative entropy penalty coefficient. |
| jump_failure_budget: Maximum cumulative failure risk for forced jumps. |
| vocab_chunk_size: Vocabulary projection planning size recorded with the |
| model. Inference uses standard PyTorch tensor operations. |
| latent_dim: Width of the recurrent latent state. |
| latent_memory_slots: Number of persistent latent memory slots. |
| latent_num_layers: Number of latent Transformer blocks. |
| latent_num_heads: Number of latent attention heads. |
| latent_local_attention_window: Local token-attention radius. |
| latent_dropout: Latent Transformer dropout probability. |
| jump_on_no_progress_after: Stagnation steps before a forced jump. |
| max_ponder_steps: Maximum denoising iterations per requested token. |
| min_trajectory_progress: Minimum fused-risk improvement counted as progress. |
| turn_end_token_id: Native Gemma turn terminator. |
| kwargs: Standard DiffusionGemma configuration values. |
| """ |
|
|
| model_type = "modilify_mk1" |
| sub_configs = { |
| "text_config": ModilifyMk1TextConfig, |
| **{ |
| key: value |
| for key, value in DiffusionGemmaConfig.sub_configs.items() |
| if key != "text_config" |
| }, |
| } |
|
|
| def __init__( |
| self, |
| text_config: ( |
| ModilifyMk1TextConfig |
| | DiffusionGemmaTextConfig |
| | dict[str, Any] |
| | None |
| ) = None, |
| vision_config: Any | dict[str, Any] | None = None, |
| *, |
| denoise_temperature: float = 0.8, |
| commit_failure_budget: float = 0.2, |
| fused_entropy_weight: float = 0.5, |
| jump_failure_budget: float = 2.0, |
| vocab_chunk_size: int = 65_536, |
| latent_dim: int = 1536, |
| latent_memory_slots: int = 64, |
| latent_num_layers: int = 4, |
| latent_num_heads: int = 16, |
| latent_local_attention_window: int = 128, |
| latent_dropout: float = 0.0, |
| jump_on_no_progress_after: int = 12, |
| max_ponder_steps: int = 64, |
| min_trajectory_progress: float = 0.005, |
| turn_end_token_id: int = 106, |
| **kwargs: Any, |
| ) -> None: |
| kwargs.pop("model_type", None) |
| if isinstance(text_config, DiffusionGemmaTextConfig): |
| text_payload = text_config.to_dict() |
| text_payload.pop("model_type", None) |
| text_config = ModilifyMk1TextConfig(**text_payload) |
| elif isinstance(text_config, dict): |
| text_payload = dict(text_config) |
| text_payload.pop("model_type", None) |
| text_config = ModilifyMk1TextConfig(**text_payload) |
| elif text_config is None: |
| text_config = ModilifyMk1TextConfig() |
|
|
| self.denoise_temperature = float(denoise_temperature) |
| self.commit_failure_budget = float(commit_failure_budget) |
| self.fused_entropy_weight = float(fused_entropy_weight) |
| self.jump_failure_budget = float(jump_failure_budget) |
| self.vocab_chunk_size = int(vocab_chunk_size) |
| self.latent_dim = int(latent_dim) |
| self.latent_memory_slots = int(latent_memory_slots) |
| self.latent_num_layers = int(latent_num_layers) |
| self.latent_num_heads = int(latent_num_heads) |
| self.latent_local_attention_window = int(latent_local_attention_window) |
| self.latent_dropout = float(latent_dropout) |
| self.jump_on_no_progress_after = int(jump_on_no_progress_after) |
| self.max_ponder_steps = int(max_ponder_steps) |
| self.min_trajectory_progress = float(min_trajectory_progress) |
| self.turn_end_token_id = int(turn_end_token_id) |
| super().__init__( |
| text_config=text_config, |
| vision_config=vision_config, |
| **kwargs, |
| ) |
| self.model_type = type(self).model_type |
| if not hasattr(self, "eos_token_id"): |
| self.eos_token_id = self.text_config.eos_token_id |
| if not hasattr(self, "pad_token_id"): |
| self.pad_token_id = self.text_config.pad_token_id |
| if not hasattr(self, "bos_token_id"): |
| self.bos_token_id = self.text_config.bos_token_id |
| self._validate_modilify() |
|
|
| def _validate_modilify(self) -> None: |
| """Validate inference-only extension values.""" |
|
|
| policy_values = ( |
| self.denoise_temperature, |
| self.commit_failure_budget, |
| self.fused_entropy_weight, |
| self.jump_failure_budget, |
| self.min_trajectory_progress, |
| ) |
| if any(not math.isfinite(value) for value in policy_values): |
| raise ValueError("Modilify Mk1 policy values must be finite.") |
| positive = ( |
| self.denoise_temperature, |
| self.commit_failure_budget, |
| self.jump_failure_budget, |
| self.vocab_chunk_size, |
| self.latent_dim, |
| self.latent_memory_slots, |
| self.latent_num_layers, |
| self.latent_num_heads, |
| self.latent_local_attention_window, |
| self.jump_on_no_progress_after, |
| self.max_ponder_steps, |
| ) |
| if any(value <= 0 for value in positive): |
| raise ValueError( |
| "Modilify Mk1 dimensions, budgets, and intervals must be positive." |
| ) |
| if self.fused_entropy_weight < 0: |
| raise ValueError("`fused_entropy_weight` must be non-negative.") |
| if self.latent_dim % self.latent_num_heads: |
| raise ValueError("`latent_dim` must be divisible by `latent_num_heads`.") |
| if not 0.0 <= self.latent_dropout < 1.0: |
| raise ValueError("`latent_dropout` must be in [0, 1).") |
| if self.min_trajectory_progress < 0: |
| raise ValueError("`min_trajectory_progress` must be non-negative.") |
|
|
|
|
| __all__ = ["ModilifyMk1Config", "ModilifyMk1TextConfig"] |
|
|
