| """Text-only schema25 inference configuration.""" |
| from __future__ import annotations |
| from collections.abc import Sequence |
| from typing import Any |
| from transformers import PreTrainedConfig |
| from transformers.models.diffusion_gemma import DiffusionGemmaTextConfig |
|
|
| DENOISE_TEMPERATURE = 0.8 |
|
|
| class ModilifyMk2TextConfig(DiffusionGemmaTextConfig): |
| model_type = "modilify_mk2_text" |
| vocab_size: int = 262_144 |
| hidden_size: int = 2816 |
| intermediate_size: int = 2112 |
| num_hidden_layers: int = 30 |
| num_attention_heads: int = 16 |
| num_key_value_heads: int = 8 |
| head_dim: int = 256 |
| max_position_embeddings: int = 262_144 |
| sliding_window: int = 1024 |
| use_bidirectional_attention: str | None = None |
| num_global_key_value_heads: int | None = 2 |
| global_head_dim: int = 512 |
| num_experts: int | None = 128 |
| top_k_experts: int | None = 8 |
| moe_intermediate_size: int | None = 704 |
|
|
| class ModilifyMk2Config(PreTrainedConfig): |
| model_type = "modilify_mk2" |
| sub_configs = {"text_config": ModilifyMk2TextConfig} |
|
|
| def __init__( |
| self, text_config: ModilifyMk2TextConfig | dict[str, Any] | None = None, *, |
| canvas_length: int = 256, |
| initializer_range: float = 0.02, |
| tie_word_embeddings: bool = True, |
| state_schema_version: int = 25, |
| memory_architecture: str = "compact_gdn2_v2", |
| latent_dim: int = 2816, |
| latent_ffn_dim: int = 7168, |
| latent_num_layers: int = 4, |
| latent_num_heads: int = 16, |
| latent_local_attention_window: int = 128, |
| latent_tape_probes: int = 4, |
| latent_history_kv_rank: int = 1024, |
| latent_working_last_block_global: bool = True, |
| working_memory_bus: bool = True, |
| persistent_memory_bus: bool = True, |
| commit_sequence_dim: int = 1024, |
| kv_cache_bucket_size: int = 128, |
| vocab_chunk_size: int = 32768, |
| turn_end_token_id: int = 106, |
| terminal_token_ids: Sequence[int] = (106, 50), |
| eos_token_id: int = 1, |
| commit_failure_budget: float = 0.2, |
| commit_top_k: int | None = 40, |
| commit_min_p: float | None = 0.05, |
| commit_target_confidence: float | None = 0.5, |
| commit_entropy_weight: float = 1.0, |
| commit_confidence_power: float = 1.0, |
| **kwargs: Any, |
| ) -> None: |
| if state_schema_version != 25 or memory_architecture != "compact_gdn2_v2": |
| raise ValueError("This release requires schema25 compact_gdn2_v2 weights.") |
| if text_config is None: |
| text_config = ModilifyMk2TextConfig() |
| elif isinstance(text_config, dict): |
| payload = dict(text_config) |
| payload.pop("model_type", None) |
| text_config = ModilifyMk2TextConfig(**payload) |
| text_config.use_bidirectional_attention = None |
| self.text_config = text_config |
| self.canvas_length = canvas_length |
| self.initializer_range = initializer_range |
| self.state_schema_version = state_schema_version |
| self.memory_architecture = memory_architecture |
| self.latent_dim = latent_dim |
| self.latent_ffn_dim = latent_ffn_dim |
| self.latent_num_layers = latent_num_layers |
| self.latent_num_heads = latent_num_heads |
| self.latent_local_attention_window = latent_local_attention_window |
| self.latent_tape_probes = latent_tape_probes |
| self.latent_history_kv_rank = latent_history_kv_rank |
| self.latent_working_last_block_global = latent_working_last_block_global |
| self.working_memory_bus = working_memory_bus |
| self.persistent_memory_bus = persistent_memory_bus |
| self.commit_sequence_dim = commit_sequence_dim |
| self.kv_cache_bucket_size = kv_cache_bucket_size |
| self.vocab_chunk_size = vocab_chunk_size |
| self.turn_end_token_id = turn_end_token_id |
| self.terminal_token_ids = terminal_token_ids |
| self.commit_failure_budget = commit_failure_budget |
| self.commit_top_k = commit_top_k |
| self.commit_min_p = commit_min_p |
| self.commit_target_confidence = commit_target_confidence |
| self.commit_entropy_weight = commit_entropy_weight |
| self.commit_confidence_power = commit_confidence_power |
| super().__init__(tie_word_embeddings=tie_word_embeddings, |
| eos_token_id=eos_token_id, **kwargs) |
|
|
