File size: 6,937 Bytes
164d101 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | # Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""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"]
|