theme_generation/engines/transformer.py

"""Small decoder-only transformer engine for symbolic theme generation."""

from __future__ import annotations

import math
import random
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from ..common import START_SYMBOL, Symbol
from ..constraints import make_theme_acceptor


PAD_SYMBOL = "__PAD__"


def require_torch():
    try:
        import torch
        from torch import nn
    except ImportError as exc:  # pragma: no cover - depends on optional dependency.
        raise RuntimeError(
            "The transformer engine requires PyTorch. Install torch in the Python "
            "environment used to run scripts/run_transformer_baseline.py."
        ) from exc
    return torch, nn


@dataclass(frozen=True)
class TransformerConfig:
    block_size: int = 64
    d_model: int = 96
    nhead: int = 4
    num_layers: int = 3
    dim_feedforward: int = 192
    dropout: float = 0.1
    batch_size: int = 64
    steps: int = 800
    learning_rate: float = 3e-4
    temperature: float = 1.0
    top_k: int = 16
    max_retries: int = 100


class SymbolVocab:
    def __init__(self, symbols: list[Symbol]) -> None:
        self.tokens: list[object] = [PAD_SYMBOL, START_SYMBOL] + symbols
        self.index = {token: idx for idx, token in enumerate(self.tokens)}

    @property
    def pad_id(self) -> int:
        return self.index[PAD_SYMBOL]

    @property
    def start_id(self) -> int:
        return self.index[START_SYMBOL]

    def encode_symbol(self, symbol: Symbol) -> int:
        return self.index[symbol]

    def decode_symbol(self, token_id: int) -> Symbol:
        token = self.tokens[token_id]
        if not isinstance(token, Symbol):
            raise ValueError(f"Token {token!r} is not a musical symbol")
        return token

    def __len__(self) -> int:
        return len(self.tokens)


@dataclass(frozen=True)
class TransformerCheckpoint:
    model: Any
    vocab: SymbolVocab
    cfg: TransformerConfig
    diagnostics: dict[str, Any]
    device: str


def make_vocab(sequences: list[tuple[Symbol, ...]]) -> SymbolVocab:
    symbols = sorted({symbol for seq in sequences for symbol in seq}, key=lambda s: (s.rpc, s.duration))
    return SymbolVocab(symbols)


def vocab_to_payload(vocab: SymbolVocab) -> list[tuple[int, str]]:
    return [(token.rpc, token.duration) for token in vocab.tokens if isinstance(token, Symbol)]


def vocab_from_payload(payload: list[tuple[int, str]] | list[list[Any]]) -> SymbolVocab:
    return SymbolVocab([Symbol(int(rpc), str(duration)) for rpc, duration in payload])


def make_training_data(sequences: list[tuple[Symbol, ...]], vocab: SymbolVocab, block_size: int) -> list[list[int]]:
    examples = []
    for seq in sequences:
        ids = [vocab.start_id] + [vocab.encode_symbol(symbol) for symbol in seq]
        if len(ids) <= block_size + 1:
            examples.append(ids)
            continue
        for start in range(0, len(ids) - block_size, block_size):
            examples.append(ids[start : start + block_size + 1])
    return examples


def batch_examples(torch, examples: list[list[int]], *, batch_size: int, block_size: int, pad_id: int, rng: random.Random):
    batch = [examples[rng.randrange(len(examples))] for _ in range(batch_size)]
    x_rows = []
    y_rows = []
    for ids in batch:
        x = ids[:-1][:block_size]
        y = ids[1:][:block_size]
        if len(x) < block_size:
            x = x + [pad_id] * (block_size - len(x))
            y = y + [pad_id] * (block_size - len(y))
        x_rows.append(x)
        y_rows.append(y)
    return torch.tensor(x_rows, dtype=torch.long), torch.tensor(y_rows, dtype=torch.long)


def build_model(nn, vocab_size: int, cfg: TransformerConfig):
    class TinyThemeTransformer(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.token_embedding = nn.Embedding(vocab_size, cfg.d_model)
            self.position_embedding = nn.Embedding(cfg.block_size, cfg.d_model)
            layer = nn.TransformerEncoderLayer(
                d_model=cfg.d_model,
                nhead=cfg.nhead,
                dim_feedforward=cfg.dim_feedforward,
                dropout=cfg.dropout,
                batch_first=True,
                norm_first=True,
            )
            self.blocks = nn.TransformerEncoder(layer, num_layers=cfg.num_layers)
            self.norm = nn.LayerNorm(cfg.d_model)
            self.head = nn.Linear(cfg.d_model, vocab_size)

        def forward(self, idx):
            batch_size, seq_len = idx.shape
            pos = idx.new_tensor(range(seq_len)).unsqueeze(0).expand(batch_size, seq_len)
            x = self.token_embedding(idx) + self.position_embedding(pos)
            mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=idx.device)
            x = self.blocks(x, mask=mask, is_causal=True)
            return self.head(self.norm(x))

    return TinyThemeTransformer()


def choose_device(torch, requested: str) -> str:
    if requested != "auto":
        return requested
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def train_model(
    *,
    sequences: list[tuple[Symbol, ...]],
    cfg: TransformerConfig,
    seed: int,
    requested_device: str,
) -> tuple[Any, SymbolVocab, dict[str, Any]]:
    torch, nn = require_torch()
    torch.manual_seed(seed)
    rng = random.Random(seed)
    vocab = make_vocab(sequences)
    examples = make_training_data(sequences, vocab, cfg.block_size)
    device = choose_device(torch, requested_device)

    model = build_model(nn, len(vocab), cfg).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.learning_rate)
    loss_fn = nn.CrossEntropyLoss(ignore_index=vocab.pad_id)

    model.train()
    last_loss = math.nan
    for _ in range(cfg.steps):
        x, y = batch_examples(
            torch,
            examples,
            batch_size=cfg.batch_size,
            block_size=cfg.block_size,
            pad_id=vocab.pad_id,
            rng=rng,
        )
        x = x.to(device)
        y = y.to(device)
        optimizer.zero_grad(set_to_none=True)
        logits = model(x)
        loss = loss_fn(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
        loss.backward()
        optimizer.step()
        last_loss = float(loss.detach().cpu())

    diagnostics = {
        "engine": "tiny decoder-only transformer",
        "training examples": len(examples),
        "training steps": cfg.steps,
        "final train loss": f"{last_loss:.4f}",
        "device": device,
        "parameters": sum(param.numel() for param in model.parameters()),
    }
    return model, vocab, diagnostics


def save_transformer_checkpoint(
    *,
    path: Path,
    model,
    vocab: SymbolVocab,
    cfg: TransformerConfig,
    diagnostics: dict[str, Any],
) -> None:
    torch, _ = require_torch()
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "format": "theme-transformer-checkpoint-v1",
            "cfg": asdict(cfg),
            "vocab": vocab_to_payload(vocab),
            "diagnostics": diagnostics,
            "state_dict": model.state_dict(),
        },
        path,
    )


def train_and_save_checkpoint(
    *,
    sequences: list[tuple[Symbol, ...]],
    cfg: TransformerConfig,
    seed: int,
    requested_device: str,
    path: Path,
) -> TransformerCheckpoint:
    model, vocab, diagnostics = train_model(
        sequences=sequences,
        cfg=cfg,
        seed=seed,
        requested_device=requested_device,
    )
    save_transformer_checkpoint(path=path, model=model, vocab=vocab, cfg=cfg, diagnostics=diagnostics)
    return TransformerCheckpoint(
        model=model,
        vocab=vocab,
        cfg=cfg,
        diagnostics=diagnostics,
        device=diagnostics["device"],
    )


def load_transformer_checkpoint(path: Path, *, requested_device: str = "auto") -> TransformerCheckpoint:
    torch, nn = require_torch()
    device = choose_device(torch, requested_device)
    payload = torch.load(path, map_location=device)
    if payload.get("format") != "theme-transformer-checkpoint-v1":
        raise RuntimeError(f"Unsupported transformer checkpoint format in {path}")
    cfg = TransformerConfig(**payload["cfg"])
    vocab = vocab_from_payload(payload["vocab"])
    model = build_model(nn, len(vocab), cfg).to(device)
    model.load_state_dict(payload["state_dict"])
    model.eval()
    diagnostics = dict(payload.get("diagnostics", {}))
    diagnostics["device"] = device
    diagnostics["checkpoint"] = str(path)
    diagnostics["engine"] = "tiny decoder-only transformer checkpoint"
    return TransformerCheckpoint(model=model, vocab=vocab, cfg=cfg, diagnostics=diagnostics, device=device)


def sample_next_id(
    torch,
    logits,
    *,
    allowed_ids: list[int],
    allowed_weights: list[float],
    temperature: float,
    top_k: int,
    rng: random.Random,
) -> int:
    values = logits[allowed_ids].detach().float().cpu()
    weight_log = torch.tensor([math.log(max(weight, 1e-12)) for weight in allowed_weights], dtype=values.dtype)
    values = values + weight_log
    if top_k > 0 and len(allowed_ids) > top_k:
        top_values, top_positions = torch.topk(values, k=top_k)
        allowed_ids = [allowed_ids[int(pos)] for pos in top_positions]
        values = top_values
    if temperature <= 0:
        return allowed_ids[int(torch.argmax(values))]
    probs = torch.softmax(values / temperature, dim=0).tolist()
    return rng.choices(allowed_ids, weights=probs, k=1)[0]


def generate_one(
    *,
    torch,
    model,
    vocab: SymbolVocab,
    length: int,
    cfg: TransformerConfig,
    start_weights: dict[int, float],
    end_weights: dict[int, float],
    endpoint_strength: float,
    enforce_triplet_groups: bool,
    rng: random.Random,
) -> tuple[Symbol, ...] | None:
    alphabet = [token for token in vocab.tokens if isinstance(token, Symbol)]
    acceptor = make_theme_acceptor(
        length=length,
        alphabet=alphabet,
        start_weights=start_weights,
        end_weights=end_weights,
        strength=endpoint_strength,
        enforce_triplet_groups=enforce_triplet_groups,
    )
    device = next(model.parameters()).device
    state = acceptor.start_state
    ids = [vocab.start_id]
    sequence = []
    model.eval()
    with torch.no_grad():
        for _ in range(length):
            allowed_ids = []
            allowed_weights = []
            for symbol in alphabet:
                next_state = acceptor.next_state(state, symbol)
                if next_state is not None:
                    allowed_ids.append(vocab.encode_symbol(symbol))
                    allowed_weights.append(acceptor.transition_weight(state, symbol))
            if not allowed_ids:
                return None
            context = ids[-cfg.block_size :]
            x = torch.tensor([context], dtype=torch.long, device=device)
            logits = model(x)[0, -1]
            chosen_id = sample_next_id(
                torch,
                logits,
                allowed_ids=allowed_ids,
                allowed_weights=allowed_weights,
                temperature=cfg.temperature,
                top_k=cfg.top_k,
                rng=rng,
            )
            symbol = vocab.decode_symbol(chosen_id)
            next_state = acceptor.next_state(state, symbol)
            if next_state is None:
                return None
            ids.append(chosen_id)
            sequence.append(symbol)
            state = next_state
    if not acceptor.is_accepting(state):
        return None
    return tuple(sequence)


def generate_transformer(
    *,
    sequences: list[tuple[Symbol, ...]],
    length: int,
    samples: int,
    start_weights: dict[int, float],
    end_weights: dict[int, float],
    endpoint_strength: float,
    enforce_triplet_groups: bool,
    seed: int,
    cfg: TransformerConfig,
    device: str,
) -> tuple[list[tuple[Symbol, ...]], dict[str, Any]]:
    torch, _ = require_torch()
    model, vocab, diagnostics = train_model(
        sequences=sequences,
        cfg=cfg,
        seed=seed,
        requested_device=device,
    )
    rng = random.Random(seed + 1)
    generated = []
    attempts = 0
    while len(generated) < samples and attempts < cfg.max_retries:
        attempts += 1
        sequence = generate_one(
            torch=torch,
            model=model,
            vocab=vocab,
            length=length,
            cfg=cfg,
            start_weights=start_weights,
            end_weights=end_weights,
            endpoint_strength=endpoint_strength,
            enforce_triplet_groups=enforce_triplet_groups,
            rng=rng,
        )
        if sequence is not None:
            generated.append(sequence)
    if len(generated) < samples:
        raise RuntimeError(f"Only generated {len(generated)} accepted samples after {attempts} attempts")
    diagnostics["accepted attempts"] = attempts
    diagnostics["vocabulary size"] = len(vocab)
    return generated, diagnostics


def sample_transformer_checkpoint(
    *,
    checkpoint: TransformerCheckpoint,
    length: int,
    samples: int,
    start_weights: dict[int, float],
    end_weights: dict[int, float],
    endpoint_strength: float,
    enforce_triplet_groups: bool,
    seed: int,
    temperature: float | None = None,
    top_k: int | None = None,
    max_retries: int | None = None,
) -> tuple[list[tuple[Symbol, ...]], dict[str, Any]]:
    torch, _ = require_torch()
    cfg = TransformerConfig(
        **{
            **asdict(checkpoint.cfg),
            "temperature": checkpoint.cfg.temperature if temperature is None else temperature,
            "top_k": checkpoint.cfg.top_k if top_k is None else top_k,
            "max_retries": checkpoint.cfg.max_retries if max_retries is None else max_retries,
        }
    )
    rng = random.Random(seed + 1)
    generated = []
    attempts = 0
    while len(generated) < samples and attempts < cfg.max_retries:
        attempts += 1
        sequence = generate_one(
            torch=torch,
            model=checkpoint.model,
            vocab=checkpoint.vocab,
            length=length,
            cfg=cfg,
            start_weights=start_weights,
            end_weights=end_weights,
            endpoint_strength=endpoint_strength,
            enforce_triplet_groups=enforce_triplet_groups,
            rng=rng,
        )
        if sequence is not None:
            generated.append(sequence)
    if len(generated) < samples:
        raise RuntimeError(f"Only generated {len(generated)} accepted samples after {attempts} attempts")
    diagnostics = dict(checkpoint.diagnostics)
    diagnostics["accepted attempts"] = attempts
    diagnostics["vocabulary size"] = len(checkpoint.vocab)
    diagnostics["temperature"] = cfg.temperature
    diagnostics["top k"] = cfg.top_k
    return generated, diagnostics