logit_conditioning_pseudocode.py

"""Pseudocode for Matrix-glyph logit conditioning in a diffusion sampler.

Use this in a runtime that exposes per-step logits before sampling. The raw
Transformers DiffusionGemma path is not recommended for a 32 GB Mac because the
unquantized model needs far more memory than a local laptop usually has.
"""

from __future__ import annotations

import torch


def token_ids_matching_unicode_blocks(tokenizer, blocks: tuple[tuple[int, int], ...]) -> list[int]:
    """Return token ids whose decoded form consists only of chars in `blocks`."""
    keep: list[int] = []
    vocab_size = len(tokenizer)
    for token_id in range(vocab_size):
        piece = tokenizer.decode([token_id], skip_special_tokens=True)
        if not piece:
            continue
        if all(any(start <= ord(char) <= end for start, end in blocks) for char in piece):
            keep.append(token_id)
    return keep


def make_matrix_prior(
    logits: torch.Tensor,
    matrix_token_ids: list[int],
    boost: float = 18.0,
    floor: float = -30.0,
) -> torch.Tensor:
    """Create logits that strongly prefer Matrix-looking glyph tokens."""
    prior = torch.full_like(logits, floor)
    prior[..., matrix_token_ids] = boost
    return prior


def alpha_schedule(step: int, total_steps: int, hard_mask_fraction: float = 0.22) -> float:
    """Start with a hard Matrix phase, then ease into model logits."""
    progress = step / max(1, total_steps - 1)
    if progress < hard_mask_fraction:
        return 0.0
    scaled = (progress - hard_mask_fraction) / (1.0 - hard_mask_fraction)
    return scaled * scaled * (3.0 - 2.0 * scaled)


def condition_logits(
    model_logits: torch.Tensor,
    matrix_token_ids: list[int],
    step: int,
    total_steps: int,
    temperature: float = 1.0,
) -> torch.Tensor:
    """Blend Matrix-glyph prior with model logits for one denoising step."""
    alpha = alpha_schedule(step, total_steps)
    matrix_logits = make_matrix_prior(model_logits, matrix_token_ids)
    blended = (1.0 - alpha) * matrix_logits + alpha * model_logits
    return blended / max(temperature, 1e-6)


def sampler_loop_pseudocode(model, tokenizer, initial_canvas, total_steps: int = 32):
    """Sketch only: adapt names to the actual DiffusionGemma sampler internals."""
    kana_blocks = (
        (0x30A0, 0x30FF),  # Katakana
        (0xFF65, 0xFF9F),  # Halfwidth Katakana
        (0x3040, 0x309F),  # Hiragana
    )
    matrix_token_ids = token_ids_matching_unicode_blocks(tokenizer, kana_blocks)
    canvas = initial_canvas

    for step in range(total_steps):
        model_logits = model.forward_diffusion_step(canvas, step=step).logits
        logits = condition_logits(model_logits, matrix_token_ids, step, total_steps)
        canvas = torch.distributions.Categorical(logits=logits).sample()
        print(tokenizer.decode(canvas[0], skip_special_tokens=True))

    return canvas