generation_modilify_mk1.py

# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Rolling generation for latent-memory Modilify Mk1."""

from __future__ import annotations

from dataclasses import dataclass, replace
import math
from typing import Any

import torch
from transformers.cache_utils import Cache
from transformers.generation import LogitsProcessorList
from transformers.generation.streamers import BaseStreamer
from transformers.modeling_outputs import ModelOutput

from transformers.models.diffusion_gemma import (
    DiffusionGemmaGenerationConfig,
    DiffusionGemmaGenerationMixin,
)
from .commit_policy import fused_commit_failure_rate, select_commit_lengths
from .latent_deliberation import LatentDeliberationState


class ModilifyMk1GenerationConfig(DiffusionGemmaGenerationConfig):
    """Generation controls for the Modilify Mk1 commit policy.

    Args:
        turn_end_token_id: Token that closes a native Gemma turn.
        denoise_temperature: Sampling temperature used for every canvas step.
        kwargs: Standard DiffusionGemma generation arguments.
    """

    def __init__(
        self,
        *,
        turn_end_token_id: int | None = None,
        denoise_temperature: float = 0.8,
        **kwargs: Any,
    ) -> None:
        self.turn_end_token_id = turn_end_token_id
        self.denoise_temperature = float(denoise_temperature)
        kwargs.pop("one_token_per_denoise_step", None)
        kwargs["t_min"] = self.denoise_temperature
        kwargs["t_max"] = self.denoise_temperature
        super().__init__(**kwargs)
        self.one_token_per_denoise_step = False

    def update(self, **kwargs: Any) -> dict[str, Any]:
        """Apply standard generation overrides and one temperature override."""

        if "denoise_temperature" in kwargs:
            self.denoise_temperature = float(kwargs.pop("denoise_temperature"))
        kwargs["t_min"] = self.denoise_temperature
        kwargs["t_max"] = self.denoise_temperature
        unused = super().update(**kwargs)
        self.one_token_per_denoise_step = False
        self.t_min = self.denoise_temperature
        self.t_max = self.denoise_temperature
        return unused

    def validate(self, **kwargs: Any) -> None:
        """Validate fixed-temperature generation values.

        DiffusionGemma's parent validator requires a non-empty temperature
        interval. Modilify Mk1 intentionally uses one fixed temperature, so the
        equivalent ``t_min == t_max`` configuration is validated here.
        """

        del kwargs
        if (
            not math.isfinite(self.denoise_temperature)
            or self.denoise_temperature <= 0
        ):
            raise ValueError("`denoise_temperature` must be positive.")
        if self.max_denoising_steps is not None and (
            not isinstance(self.max_denoising_steps, int)
            or self.max_denoising_steps <= 0
        ):
            raise ValueError("`max_denoising_steps` must be a positive integer.")
        if self.turn_end_token_id is not None and (
            not isinstance(self.turn_end_token_id, int) or self.turn_end_token_id < 0
        ):
            raise ValueError("`turn_end_token_id` must be a non-negative integer.")

    @classmethod
    def from_model_config(cls, model_config: Any) -> "ModilifyMk1GenerationConfig":
        """Build generation defaults from a model configuration."""

        return cls(
            turn_end_token_id=model_config.turn_end_token_id,
            denoise_temperature=model_config.denoise_temperature,
            eos_token_id=getattr(
                model_config,
                "eos_token_id",
                model_config.text_config.eos_token_id,
            ),
        )

    @staticmethod
    def _get_default_generation_params() -> dict[str, object]:
        """Return defaults with no inherited entropy/readiness commit controls."""

        return {
            "max_new_tokens": 256,
            "max_denoising_steps": 48,
            "t_min": 0.8,
            "t_max": 0.8,
        }


@dataclass
class ModilifyMk1GenerationOutput(ModelOutput):
    """Structured result returned by rolling block-diffusion generation."""

    sequences: torch.LongTensor
    generated_lengths: torch.LongTensor | None = None
    tokens_per_forward: torch.FloatTensor | None = None
    past_key_values: Cache | None = None
    stop_reason: str | tuple[str, ...] | None = None
    committed_tokens: int | torch.LongTensor | None = None
    denoise_steps: int | torch.LongTensor | None = None
    no_progress_steps: int | torch.LongTensor | None = None
    jump_count: int | torch.LongTensor | None = None
    forced_jump_bad_count: int | torch.LongTensor | None = None
    heavy_forward_count: int | torch.LongTensor | None = None
    latent_context_update_count: int | torch.LongTensor | None = None
    average_commit_len: float | torch.FloatTensor | None = None
    state_shift_count: int | torch.LongTensor | None = None
    latent_memory_norm: float | torch.FloatTensor | None = None
    state_retention_score: float | torch.FloatTensor | None = None
    logits: None = None
    scores: None = None
    hidden_states: None = None


@dataclass
class _RollingState:
    """All real iterative state; no vocabulary-sized tensor is retained."""

    canvas: torch.LongTensor
    confidence: torch.FloatTensor
    entropy: torch.FloatTensor
    age: torch.IntTensor
    latent_state: LatentDeliberationState
    history_hidden_state: torch.FloatTensor | None


def _retain_denoise_proposals(proposal: torch.LongTensor) -> torch.LongTensor:
    """Keep every latest denoise token; confidence controls commit, not writeback."""
    if proposal.ndim != 2:
        raise ValueError("Denoise proposals must have shape [batch, canvas].")
    return proposal.clone()


class _NoiseCanvasSampler:
    """Uniform diffusion noise source with no commit-policy responsibilities."""

    def __init__(self, *, canvas_length: int, vocab_size: int) -> None:
        self.canvas_length = int(canvas_length)
        self.vocab_size = int(vocab_size)
        self.initial_entropy = math.log(self.vocab_size)

    def initialize_canvas(
        self,
        batch_size: int,
        device: torch.device,
    ) -> torch.LongTensor:
        """Sample a uniformly random starting canvas.

        Args:
            batch_size: Number of canvases to create.
            device: Device on which token IDs are allocated.

        Returns:
            Random token IDs with shape ``[batch_size, canvas_length]``.
        """

        return torch.randint(
            self.vocab_size,
            (batch_size, self.canvas_length),
            device=device,
        )


class ModilifyMk1GenerationMixin(DiffusionGemmaGenerationMixin):
    """Transformers-compatible rolling latent-deliberation generator."""

    def _prepare_sampler(
        self,
        generation_config: ModilifyMk1GenerationConfig,
        canvas_length: int | None = None,
    ) -> _NoiseCanvasSampler:
        del generation_config
        return _NoiseCanvasSampler(
            canvas_length=canvas_length or self.config.canvas_length,
            vocab_size=self.config.text_config.vocab_size,
        )

    @staticmethod
    def _shift_state_rows(
        state: _RollingState,
        commit_lengths: torch.LongTensor,
        sampler: _NoiseCanvasSampler,
    ) -> _RollingState:
        """Shift every rolling row by its own committed prefix length."""

        batch_size, canvas_length = state.canvas.shape
        if commit_lengths.shape != (batch_size,):
            raise ValueError("Commit lengths must have shape [batch].")
        if not bool(commit_lengths.gt(0).any()):
            return state
        positions = torch.arange(canvas_length, device=state.canvas.device)[None, :]
        source = positions + commit_lengths[:, None]
        retained = source.lt(canvas_length)

        def shift(value: torch.Tensor, fill_value: float | int = 0) -> torch.Tensor:
            index = source.clamp_max(canvas_length - 1)
            index = index.view(
                batch_size, canvas_length, *([1] * (value.ndim - 2))
            ).expand_as(value)
            gathered = value.gather(1, index)
            mask = retained.view(
                batch_size, canvas_length, *([1] * (value.ndim - 2))
            )
            fill = torch.as_tensor(fill_value, device=value.device, dtype=value.dtype)
            return torch.where(mask, gathered, fill)

        tail = sampler.initialize_canvas(batch_size, state.canvas.device)
        canvas = torch.cat((state.canvas, tail), dim=1).gather(1, source)
        unknown_entropy = float(sampler.initial_entropy)
        latent = state.latent_state
        committed = commit_lengths.gt(0)
        shifted_latent = LatentDeliberationState(
            token_latents=shift(latent.token_latents),
            memory_slots=latent.memory_slots.clone(),
            confidence=shift(latent.confidence),
            entropy=shift(latent.entropy, unknown_entropy),
            age=shift(latent.age),
            token_changed=shift(latent.token_changed),
            confidence_delta=shift(latent.confidence_delta),
            entropy_delta=shift(latent.entropy_delta),
            ponder_steps=torch.where(
                committed, torch.zeros_like(latent.ponder_steps), latent.ponder_steps
            ),
            stagnation_steps=torch.where(
                committed, torch.zeros_like(latent.stagnation_steps), latent.stagnation_steps
            ),
        )
        return _RollingState(
            canvas=canvas,
            confidence=shift(state.confidence),
            entropy=shift(state.entropy, unknown_entropy),
            age=shift(state.age),
            latent_state=shifted_latent,
            history_hidden_state=(
                None
                if state.history_hidden_state is None
                else shift(state.history_hidden_state)
            ),
        )

    @staticmethod
    def _merge_state_rows(
        previous: _RollingState,
        updated: _RollingState,
        update_mask: torch.BoolTensor,
    ) -> _RollingState:
        """Advance active rows while leaving completed rows unchanged."""

        def choose(old: torch.Tensor, new: torch.Tensor) -> torch.Tensor:
            mask = update_mask.view(
                update_mask.shape[0],
                *([1] * (old.ndim - 1)),
            )
            return torch.where(mask, new, old)

        old_latent = previous.latent_state
        new_latent = updated.latent_state
        latent = LatentDeliberationState(
            token_latents=choose(
                old_latent.token_latents,
                new_latent.token_latents,
            ),
            memory_slots=choose(old_latent.memory_slots, new_latent.memory_slots),
            confidence=choose(old_latent.confidence, new_latent.confidence),
            entropy=choose(old_latent.entropy, new_latent.entropy),
            age=choose(old_latent.age, new_latent.age),
            token_changed=choose(
                old_latent.token_changed,
                new_latent.token_changed,
            ),
            confidence_delta=choose(
                old_latent.confidence_delta,
                new_latent.confidence_delta,
            ),
            entropy_delta=choose(
                old_latent.entropy_delta,
                new_latent.entropy_delta,
            ),
            ponder_steps=choose(
                old_latent.ponder_steps,
                new_latent.ponder_steps,
            ),
            stagnation_steps=choose(
                old_latent.stagnation_steps,
                new_latent.stagnation_steps,
            ),
        )
        history = previous.history_hidden_state
        if updated.history_hidden_state is not None:
            history = (
                updated.history_hidden_state
                if history is None
                else choose(history, updated.history_hidden_state)
            )
        return _RollingState(
            canvas=choose(previous.canvas, updated.canvas),
            confidence=choose(previous.confidence, updated.confidence),
            entropy=choose(previous.entropy, updated.entropy),
            age=choose(previous.age, updated.age),
            latent_state=latent,
            history_hidden_state=history,
        )

    @torch.inference_mode()
    def generate(
        self,
        input_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        streamer: BaseStreamer | None = None,
        generation_config: ModilifyMk1GenerationConfig | None = None,
        logits_processor: LogitsProcessorList | None = None,
        **kwargs,
    ) -> ModilifyMk1GenerationOutput:
        """Generate one or more responses with rolling block diffusion.

        Args:
            input_ids: Tokenized prompts with shape ``[batch, sequence]``.
            past_key_values: Optional existing encoder cache.
            streamer: Optional standard Transformers token streamer.
            generation_config: Generation limits and token IDs.
            logits_processor: Unsupported custom logits processors.
            **kwargs: Standard multimodal encoder inputs and generation values.

        Returns:
            Generated sequences and diffusion diagnostics.

        Raises:
            ValueError: If inputs are invalid or unsupported logits processing
                is requested.
        """

        generation_config, model_kwargs = self._prepare_generation_config(
            generation_config,
            **kwargs,
        )
        if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] < 1:
            raise ValueError(
                "Modilify Mk1 generation requires `input_ids` with shape "
                "[batch, sequence]."
            )
        if logits_processor:
            raise ValueError(
                "Modilify Mk1 samples the configured fixed-temperature distribution "
                "and does not accept custom logits processors."
            )
        batch_size, input_width = input_ids.shape
        if batch_size > 1 and streamer is not None:
            raise ValueError("Streamers currently support batch size 1 only.")
        if batch_size > 1 and past_key_values is not None:
            raise ValueError("Batched generation requires a fresh KV cache.")
        device = input_ids.device
        dtype = self.model.decoder.embed_tokens.weight.dtype
        canvas_length = self.config.canvas_length
        cached_length = (
            past_key_values.get_seq_length() if past_key_values is not None else 0
        )
        _, max_new_tokens = self._prepare_generated_length(
            generation_config, cached_length + input_width
        )
        max_iterations = max(1, max_new_tokens * self.config.max_ponder_steps)
        if past_key_values is None:
            past_key_values = self._prepare_cache_for_generation(
                generation_config,
                batch_size=batch_size,
                # Ragged rows append dense masked blocks. If different rows
                # advance in different iterations, physical cache width can
                # reach the sum of all per-row generation limits.
                max_length=input_width + batch_size * max_new_tokens,
            )
        expected_mask_width = cached_length + input_width
        cache_attention_mask = model_kwargs.pop(
            "attention_mask",
            torch.ones(
                batch_size, expected_mask_width, dtype=torch.bool, device=device
            ),
        ).bool()
        if cache_attention_mask.shape != (batch_size, expected_mask_width):
            raise ValueError(
                "`attention_mask` must have shape [batch, cached_length + sequence]."
            )
        provided_position_ids = model_kwargs.pop("position_ids", None)
        if provided_position_ids is not None:
            if provided_position_ids.shape != input_ids.shape:
                raise ValueError("`position_ids` must have the same shape as `input_ids`.")
            prompt_positions = provided_position_ids.to(device=device, dtype=torch.int32)
        elif cached_length:
            prompt_positions = torch.arange(
                cached_length,
                cached_length + input_width,
                device=device,
                dtype=torch.int32,
            ).unsqueeze(0)
        else:
            input_mask = cache_attention_mask[:, -input_width:]
            prompt_positions = (
                input_mask.long()
                .cumsum(dim=-1)
                .sub(1)
                .clamp_min(0)
                .to(torch.int32)
            )
        logical_lengths = cache_attention_mask.long().sum(dim=-1)
        if input_width:
            encoder_keys = ("pixel_values", "mm_token_type_ids", "image_position_ids")
            encoder_kwargs = {
                key: model_kwargs.pop(key)
                for key in encoder_keys
                if key in model_kwargs
            }
            past_key_values = self.model.encoder(
                input_ids=input_ids,
                attention_mask=cache_attention_mask,
                past_key_values=past_key_values,
                position_ids=prompt_positions,
                **encoder_kwargs,
            ).past_key_values

        sampler = self._prepare_sampler(generation_config, canvas_length)
        latent = LatentDeliberationState.empty(
            batch_size=batch_size,
            canvas_length=canvas_length,
            latent_dim=self.config.latent_dim,
            memory_slots=self.config.latent_memory_slots,
            device=device,
            dtype=dtype,
        )
        state = _RollingState(
            canvas=sampler.initialize_canvas(batch_size, device),
            confidence=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.float32
            ),
            entropy=torch.full(
                (batch_size, canvas_length),
                math.log(self.config.text_config.vocab_size),
                device=device,
                dtype=torch.float32,
            ),
            age=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.int32
            ),
            latent_state=latent,
            history_hidden_state=None,
        )
        turn_end = (
            self.config.turn_end_token_id
            if generation_config.turn_end_token_id is None
            else generation_config.turn_end_token_id
        )
        configured_eos = generation_config.eos_token_id
        if configured_eos is None:
            configured_eos = self.config.eos_token_id
        if isinstance(configured_eos, int):
            configured_eos = [configured_eos]
        stop_token_ids = tuple(
            dict.fromkeys((int(turn_end), *(int(value) for value in configured_eos or ())))
        )
        pad_token_id = generation_config.pad_token_id
        if pad_token_id is None:
            pad_token_id = getattr(self.config, "pad_token_id", None)
        if isinstance(pad_token_id, (list, tuple)):
            pad_token_id = pad_token_id[0]
        pad_token_id = int(0 if pad_token_id is None else pad_token_id)
        generated = torch.full(
            (batch_size, max_new_tokens),
            pad_token_id,
            dtype=input_ids.dtype,
            device=device,
        )
        committed = torch.zeros(batch_size, dtype=torch.long, device=device)
        denoise_steps = torch.zeros_like(committed)
        jumps = torch.zeros_like(committed)
        forced_jump_tokens = torch.zeros_like(committed)
        shifts = torch.zeros_like(committed)
        retention_scores = torch.zeros(batch_size, dtype=torch.float32, device=device)
        stop_codes = torch.zeros_like(committed)
        active_rows = torch.ones(batch_size, dtype=torch.bool, device=device)
        canvas_positions = torch.arange(canvas_length, device=device)[None, :]
        if streamer is not None:
            streamer.put(input_ids.cpu())

        while bool(active_rows.any()):
            decoder_positions = (
                logical_lengths[:, None]
                + canvas_positions
            ).to(torch.int32)
            denoise_steps += active_rows.long()
            decoder_attention_mask = torch.cat(
                (
                    cache_attention_mask,
                    torch.ones(
                        batch_size,
                        canvas_length,
                        dtype=torch.bool,
                        device=device,
                    ),
                ),
                dim=-1,
            )
            output = self(
                input_ids=None,
                past_key_values=past_key_values,
                decoder_input_ids=state.canvas,
                previous_confidence=state.confidence,
                previous_entropy=state.entropy,
                token_age=state.age,
                latent_state=state.latent_state,
                history_hidden_state=state.history_hidden_state,
                decoder_position_ids=decoder_positions,
                decoder_read_cache=True,
                decoder_attention_mask=decoder_attention_mask,
                return_proposal_statistics=True,
                denoise_temperature=generation_config.denoise_temperature,
                **model_kwargs,
            )
            if any(
                value is None
                for value in (
                    output.proposal,
                    output.proposal_confidence,
                    output.token_entropy,
                    output.greedy_proposal,
                    output.greedy_confidence,
                )
            ):
                raise RuntimeError("Model forward did not return proposal statistics.")
            proposal = output.proposal
            proposal_confidence = output.proposal_confidence
            token_entropy = output.token_entropy
            greedy_proposal = output.greedy_proposal
            greedy_confidence = output.greedy_confidence
            next_canvas = _retain_denoise_proposals(proposal)
            next_confidence = proposal_confidence.float()
            next_latent = replace(
                output.next_latent_state,
                confidence=next_confidence.detach().float(),
                entropy=token_entropy.detach().float(),
                age=state.age + 1,
                token_changed=next_canvas.ne(state.canvas).detach().float(),
                confidence_delta=next_confidence.detach().float() - state.confidence,
                entropy_delta=token_entropy.detach().float() - state.entropy,
            )
            next_state = _RollingState(
                canvas=next_canvas,
                confidence=next_confidence,
                entropy=token_entropy,
                age=state.age + 1,
                latent_state=next_latent,
                history_hidden_state=output.heavy_hidden_state,
            )
            next_state = self._merge_state_rows(state, next_state, active_rows)
            remaining = torch.tensor(
                max_new_tokens, device=device, dtype=torch.long
            ).sub(committed)
            normal_failure_rate = fused_commit_failure_rate(
                proposal_confidence, token_entropy,
                vocab_size=self.config.text_config.vocab_size,
                entropy_weight=self.config.fused_entropy_weight,
            )
            jump_failure_rate = fused_commit_failure_rate(
                greedy_confidence, token_entropy,
                vocab_size=self.config.text_config.vocab_size,
                entropy_weight=self.config.fused_entropy_weight,
            )
            previous_failure_rate = fused_commit_failure_rate(
                state.confidence, state.entropy,
                vocab_size=self.config.text_config.vocab_size,
                entropy_weight=self.config.fused_entropy_weight,
            )
            policy_decision = select_commit_lengths(
                sampled_token_ids=proposal,
                normal_failure_rate=normal_failure_rate,
                previous_failure_rate=previous_failure_rate,
                greedy_token_ids=greedy_proposal,
                jump_failure_rate=jump_failure_rate,
                ponder_steps=state.latent_state.ponder_steps,
                stagnation_steps=state.latent_state.stagnation_steps,
                active_rows=active_rows,
                remaining_lengths=remaining,
                failure_budget=self.config.commit_failure_budget,
                jump_failure_budget=self.config.jump_failure_budget,
                stop_token_id=stop_token_ids,
                max_ponder_steps=self.config.max_ponder_steps,
                stagnation_threshold=self.config.jump_on_no_progress_after,
                min_progress=self.config.min_trajectory_progress,
            )
            next_ponder = policy_decision.ponder_steps
            next_stagnation = policy_decision.stagnation_steps
            commit_lengths = policy_decision.commit_lengths
            jump_rows = policy_decision.jump_rows
            jumps += jump_rows.long()
            forced_jump_tokens += torch.where(
                jump_rows, commit_lengths, torch.zeros_like(commit_lengths)
            )
            commit_positions = canvas_positions.lt(commit_lengths[:, None])
            if bool(jump_rows.any()):
                next_state = replace(
                    next_state,
                    canvas=torch.where(
                        commit_positions & jump_rows[:, None],
                        policy_decision.commit_token_ids,
                        next_state.canvas,
                    ),
                )
            next_state = replace(
                next_state,
                latent_state=replace(
                    next_state.latent_state,
                    ponder_steps=next_ponder,
                    stagnation_steps=next_stagnation,
                ),
            )
            commit_token_ids = policy_decision.commit_token_ids
            before = committed.clone()
            write_rows = torch.arange(batch_size, device=device)[:, None].expand_as(
                commit_token_ids
            )
            write_positions = before[:, None] + canvas_positions
            generated[
                write_rows[commit_positions], write_positions[commit_positions]
            ] = commit_token_ids[commit_positions]

            commit_width = int(commit_lengths.max())
            if commit_width:
                block_mask = torch.arange(commit_width, device=device)[None, :].lt(
                    commit_lengths[:, None]
                )
                committed_block = torch.where(
                    block_mask,
                    commit_token_ids[:, :commit_width],
                    torch.full(
                        (batch_size, commit_width),
                        pad_token_id,
                        device=device,
                        dtype=input_ids.dtype,
                    ),
                )
                block_positions = (
                    logical_lengths[:, None]
                    + canvas_positions[:, :commit_width]
                ).to(torch.int32)
                block_positions = torch.where(
                    block_mask, block_positions, torch.zeros_like(block_positions)
                )
                cache_attention_mask = torch.cat(
                    (cache_attention_mask, block_mask), dim=-1
                )
                past_key_values = self.model.encoder(
                    input_ids=committed_block,
                    attention_mask=cache_attention_mask,
                    past_key_values=past_key_values,
                    position_ids=block_positions,
                ).past_key_values
                if streamer is not None:
                    streamer.put(committed_block.cpu())
            committed += commit_lengths
            logical_lengths += commit_lengths
            committed_rows = commit_lengths.gt(0)
            shifts += committed_rows.long()
            shifted = self._shift_state_rows(next_state, commit_lengths, sampler)
            retention_scores += committed_rows.float()
            state = shifted

            turn_hits = (
                commit_token_ids.eq(turn_end) & commit_positions
            ).any(dim=-1)
            eos_hits = torch.zeros_like(turn_hits)
            for token_id in stop_token_ids:
                if token_id != turn_end:
                    eos_hits |= (
                        commit_token_ids.eq(token_id) & commit_positions
                    ).any(dim=-1)
            stop_codes = torch.where(
                stop_codes.eq(0) & turn_hits,
                torch.ones_like(stop_codes),
                stop_codes,
            )
            stop_codes = torch.where(
                stop_codes.eq(0) & eos_hits,
                torch.full_like(stop_codes, 2),
                stop_codes,
            )
            stop_codes = torch.where(
                stop_codes.eq(0) & committed.ge(max_new_tokens),
                torch.full_like(stop_codes, 3),
                stop_codes,
            )
            if generation_config.max_denoising_steps is not None:
                stop_codes = torch.where(
                    stop_codes.eq(0)
                    & denoise_steps.ge(generation_config.max_denoising_steps),
                    torch.full_like(stop_codes, 4),
                    stop_codes,
                )
            stop_codes = torch.where(
                stop_codes.eq(0) & denoise_steps.ge(max_iterations),
                torch.full_like(stop_codes, 5),
                stop_codes,
            )
            active_rows = stop_codes.eq(0)

        output_width = int(committed.max())
        sequences = torch.cat((input_ids, generated[:, :output_width]), dim=-1)
        if streamer is not None:
            streamer.end()
        reason_names = {
            1: "turn_end",
            2: "eos",
            3: "max_new_tokens",
            4: "max_denoising_steps",
            5: "episode_watchdog",
        }
        stop_reasons = tuple(
            reason_names.get(code, "unknown")
            for code in stop_codes.detach().cpu().tolist()
        )
        tokens_per_forward = committed.float() / denoise_steps.clamp_min(1).float()
        average_commit_len = committed.float() / shifts.clamp_min(1).float()
        latent_memory_norm = (
            state.latent_state.memory_slots.float().norm(dim=-1).mean(dim=-1)
        )
        state_retention_score = retention_scores / shifts.clamp_min(1).float()

        def scalar_or_tensor(
            value: torch.Tensor,
            *,
            floating: bool = False,
        ) -> int | float | torch.Tensor:
            if batch_size > 1:
                return value
            item = value[0].item()
            return float(item) if floating else int(item)

        return ModilifyMk1GenerationOutput(
            sequences=sequences,
            generated_lengths=committed.clone(),
            tokens_per_forward=tokens_per_forward,
            past_key_values=past_key_values,
            stop_reason=stop_reasons[0] if batch_size == 1 else stop_reasons,
            committed_tokens=scalar_or_tensor(committed),
            denoise_steps=scalar_or_tensor(denoise_steps),
            no_progress_steps=scalar_or_tensor(state.latent_state.stagnation_steps),
            jump_count=scalar_or_tensor(jumps),
            forced_jump_bad_count=scalar_or_tensor(forced_jump_tokens),
            heavy_forward_count=scalar_or_tensor(denoise_steps),
            latent_context_update_count=scalar_or_tensor(denoise_steps),
            average_commit_len=scalar_or_tensor(average_commit_len, floating=True),
            state_shift_count=scalar_or_tensor(shifts),
            latent_memory_norm=scalar_or_tensor(latent_memory_norm, floating=True),
            state_retention_score=scalar_or_tensor(
                state_retention_score,
                floating=True,
            ),
        )


__all__ = [
    "ModilifyMk1GenerationConfig",
    "ModilifyMk1GenerationMixin",
    "ModilifyMk1GenerationOutput",
]