Upload MaiGenerator model (epoch 10) and inference code

.gitignore

@@ -0,0 +1 @@
+__pycache__/

README.md

@@ -1,3 +1,156 @@
 ---
+language: zh
 license: mit
+tags:
+- maimai
+- music-generation
+- chart-generation
+- transformer
+- audio-conditioned
+- pytorch
+pipeline_tag: text-generation
 ---
+
+# MaiGenerator (maiChartGen)
+
+基于 EnCodec 音频编码 + Transformer 的舞萌 (maimai) 谱面自动生成模型。
+
+输入音频 + BPM/难度/等级 → 输出可玩的 maimai 谱面 (maidata.txt)。
+
+## 模型简介
+
+- **架构**: Encoder-Decoder Transformer (4 enc + 8 dec)，512 维，8 头
+- **Decoder**: 前 6 层共享 FFN + 后 2 层 MoE（4 专家 × 难度路由）
+- **核心设计**: 时间对齐 RoPE（谱面位置 = BPM 转换的音频帧号）
+- **推理**: KV-Cache 增量推理，支持 float16 加速
+- **输出**: 标准 maidata.txt 格式，可直接导入 Simai/maimai 模拟器
+
+## 快速使用
+
+### 安装依赖
+
+```bash
+pip install torch torchaudio encodec soundfile tqdm numpy
+pip install huggingface_hub
+```
+
+### 下载模型并推理
+
+```python
+from huggingface_hub import hf_hub_download
+import torch
+import subprocess
+
+# 下载模型权重
+model_path = hf_hub_download(repo_id="Goldgom/maiChartGen", filename="best.pt")
+
+# 推理
+subprocess.run([
+    "python", "inference.py",
+    "--checkpoint", model_path,
+    "--audio", "your_track.mp3",
+    "--bpm", "270", "--diff", "MASTER", "--level", "13.5",
+    "--temperature", "1.0", "--top-k", "500", "--top-p", "0.95",
+    "--density-bias", "1.0", "--target-notes", "650",
+    "--config-density-bias", "0.4", "--hold-bias", "0.2", "--slide-bias", "0.8",
+    "--break-penalty", "2.5", "--max-break-ratio", "0.06",
+    "--output", "maidata.txt"
+])
+```
+
+### 命令行推理
+
+```bash
+# 基础用法
+python inference.py --checkpoint best.pt --audio track.mp3 \
+    --bpm 150 --diff MASTER --level 12.0 --output maidata.txt
+
+# 高质量 MASTER 谱面
+python inference.py --checkpoint best.pt --audio track.mp3 \
+    --bpm 270 --diff MASTER --level 13.5 \
+    --temperature 1.0 --top-k 500 --top-p 0.95 \
+    --density-bias 1.0 --target-notes 650 \
+    --config-density-bias 0.4 --hold-bias 0.2 --slide-bias 0.8 \
+    --break-penalty 2.5 --max-break-ratio 0.06 \
+    --output maidata.txt
+
+# 高性能推理 (float16 + torch.compile)
+python inference.py --checkpoint best.pt --audio track.mp3 \
+    --bpm 150 --diff MASTER --level 12.0 \
+    --precision float16 --compile --output maidata.txt
+```
+
+### 后端 API 使用
+
+```python
+from Tokenizer.MaiTrackTokenizer import MaiTrackTokenizer
+from inference import generate_chart
+import torch
+
+# 加载模型
+device = torch.device("cuda")
+ckpt = torch.load("best.pt", map_location=device, weights_only=False)
+config = ckpt["config"]
+model = MaiGenerator(
+    d_model=config.get("d_model", 512),
+    enc_layers=config.get("enc_layers", 4),
+    dec_layers=config.get("dec_layers", 8),
+    heads=config.get("heads", 8),
+    d_ff=config.get("d_ff", 2048),
+).to(device)
+model.load_state_dict(ckpt["model_state"], strict=False)
+model.eval()
+
+# 音频 tokenize
+audio_tok = MaiTrackTokenizer(n_layers=2, device="cuda")
+audio_tokens = audio_tok.encode("track.mp3")
+
+# 生成谱面
+chart_tokens = generate_chart(
+    model, torch.tensor([audio_tokens]),
+    bpm=270, difficulty=3, level_value=13.5,
+    temperature=1.0, top_k=500, top_p=0.95,
+    density_bias=1.0, target_notes=650,
+    config_density_bias=0.4,
+    type_biases={"hold": 0.2, "slide": 0.8},
+    break_penalty=2.5, max_break_ratio=0.06,
+)
+```
+
+## 推理参数
+
+| 参数 | 默认值 | 说明 |
+|------|--------|------|
+| `--temperature` | 0.8 | 采样温度 |
+| `--top-k` | 50 | Top-K 过滤 |
+| `--top-p` | 0.95 | Nucleus 过滤 |
+| `--density-bias` | 0.0 | 密度引导强度 |
+| `--target-notes` | auto | 目标 note 数 |
+| `--hold-bias` | 0.0 | HOLD 偏置 |
+| `--slide-bias` | 0.0 | SLIDE 偏置 |
+| `--break-penalty` | 1.5 | BREAK 比例惩罚 |
+| `--max-break-ratio` | 0.08 | BREAK 比例上限 |
+| `--precision` | float32 | float16/bf16 精度 |
+
+完整参数见 [README](https://github.com/Goldgom/maiGenerator) 或 `python inference.py --help`。
+
+## 训练数据
+
+- 1,730 首歌曲，~5,000 个谱面样本
+- 难度范围: BASIC ~ ReMASTER (lv1~15)
+- 谱面类型: SD / DX / FULLTOUCH
+
+## 技术细节
+
+- **音频处理**: EnCodec 24kHz, 2-layer RVQ, 75Hz 帧率
+- **分词器**: 基础 256 token + Config Token（单 token 编码完整 note）
+- **约束解码**: 硬语法状态机 (NORMAL → SIM → DUR → SLIDE_BODY)
+- **采样引导**: 密度引导 + 类型偏置 + BREAK 比例控制
+
+## 代码仓库
+
+https://github.com/Goldgom/maiGenerator
+
+## License
+
+MIT

Tokenizer/MaiChartTokenizer.py

@@ -0,0 +1,201 @@
+"""
+MaiChartTokenizer — VQ-VAE based learned tokenizer for maimai charts.
+
+Two modes:
+  1. Rule-based: use tokenizer.MaiChartTokenizer (vocab-based, deterministic)
+  2. Learned (VQ-VAE): this module — encoder → codebook → decoder
+
+Architecture:
+  note_features (35-dim) → Encoder MLP → latent → VectorQuantizer → discrete codes
+                                                      ↓
+                              Decoder MLP → reconstructed note_features
+
+Usage:
+    from Tokenizer.MaiChartTokenizer import MaiChartTokenizer
+
+    model = MaiChartTokenizer(input_dim=35, hidden_dim=256, codebook_size=512, latent_dim=64)
+    tokens, commitment_loss = model.encode(features)   # → discrete token sequence
+    reconstructed = model.decode(tokens)               # → back to features
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VectorQuantizer(nn.Module):
+    """
+    Vector Quantization layer (VQ-VAE codebook).
+
+    Maps continuous latent vectors to the nearest codebook entry.
+    Returns discrete indices + commitment loss for training.
+    """
+
+    def __init__(self, codebook_size: int, latent_dim: int, beta: float = 0.25):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.latent_dim = latent_dim
+        self.beta = beta
+
+        # Codebook: [codebook_size, latent_dim]
+        self.embedding = nn.Embedding(codebook_size, latent_dim)
+        self.embedding.weight.data.uniform_(-1.0 / codebook_size, 1.0 / codebook_size)
+
+    def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            z: [B, T, latent_dim] or [B*T, latent_dim]
+
+        Returns:
+            z_q:     Quantized vectors (same shape as z)
+            indices: Codebook indices [B, T] or [B*T]
+            loss:    Commitment loss (scalar)
+        """
+        # Flatten if needed
+        original_shape = z.shape
+        if z.dim() == 3:
+            B, T, D = z.shape
+            z_flat = z.reshape(-1, D)
+        else:
+            z_flat = z
+
+        # Compute distances to codebook: [N, codebook_size]
+        # ||z - e||^2 = ||z||^2 + ||e||^2 - 2 * z @ e.T
+        z_sq = (z_flat ** 2).sum(dim=1, keepdim=True)
+        e_sq = (self.embedding.weight ** 2).sum(dim=1)
+        dist = z_sq + e_sq.unsqueeze(0) - 2 * z_flat @ self.embedding.weight.T
+
+        # Nearest codebook entry
+        indices = dist.argmin(dim=1)  # [N]
+        z_q = self.embedding(indices)  # [N, latent_dim]
+
+        # VQ loss: commitment loss + codebook loss
+        commitment_loss = F.mse_loss(z_flat.detach(), z_q)
+        codebook_loss = F.mse_loss(z_flat, z_q.detach())
+        vq_loss = codebook_loss + self.beta * commitment_loss
+
+        # Straight-through estimator
+        z_q = z_flat + (z_q - z_flat).detach()
+
+        # Reshape back
+        if len(original_shape) == 3:
+            z_q = z_q.reshape(B, T, D)
+            indices = indices.reshape(B, T)
+
+        return z_q, indices, vq_loss
+
+    def decode_indices(self, indices: torch.Tensor) -> torch.Tensor:
+        """Convert discrete indices back to vectors."""
+        return self.embedding(indices)
+
+
+class MaiChartTokenizer(nn.Module):
+    """
+    VQ-VAE Tokenizer: encodes chart note features → discrete token sequence.
+
+    The rule-based tokenizer (tokenizer.py) provides the baseline vocabulary.
+    This learned tokenizer can be trained end-to-end to discover optimal tokens.
+
+    Args:
+        input_dim:    Feature dimension per note (from mai_parser note_to_vector).
+        hidden_dim:   Hidden dimension of encoder/decoder MLP.
+        codebook_size: Number of discrete codes (vocabulary size).
+        latent_dim:   Dimension of each codebook vector.
+        num_layers:   Number of MLP layers in encoder/decoder.
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 35,
+        hidden_dim: int = 256,
+        codebook_size: int = 512,
+        latent_dim: int = 64,
+        num_layers: int = 2,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.codebook_size = codebook_size
+        self.latent_dim = latent_dim
+
+        # ── Encoder: note features → latent ──
+        encoder_layers = []
+        in_dim = input_dim
+        for i in range(num_layers):
+            out_dim = hidden_dim if i < num_layers - 1 else latent_dim
+            encoder_layers.append(nn.Linear(in_dim, out_dim))
+            if i < num_layers - 1:
+                encoder_layers.append(nn.ReLU())
+            in_dim = out_dim
+        self.encoder = nn.Sequential(*encoder_layers)
+
+        # ── Pre-quantization projection ──
+        self.pre_vq = nn.Linear(latent_dim, latent_dim)
+
+        # ── Vector Quantizer ──
+        self.vq = VectorQuantizer(codebook_size, latent_dim)
+
+        # ── Decoder: latent → note features ──
+        decoder_layers = []
+        in_dim = latent_dim
+        for i in range(num_layers):
+            out_dim = hidden_dim if i < num_layers - 1 else input_dim
+            decoder_layers.append(nn.Linear(in_dim, out_dim))
+            if i < num_layers - 1:
+                decoder_layers.append(nn.ReLU())
+            in_dim = out_dim
+        self.decoder = nn.Sequential(*decoder_layers)
+
+    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Encode chart features to discrete token sequence.
+
+        Args:
+            x: [B, T, input_dim] — batch of chart feature sequences.
+
+        Returns:
+            tokens: [B, T] — discrete token IDs.
+            vq_loss: scalar commitment loss (for training).
+        """
+        # Encode
+        z = self.encoder(x)            # [B, T, latent_dim]
+        z = self.pre_vq(z)
+
+        # Quantize
+        z_q, indices, vq_loss = self.vq(z)
+
+        return indices, vq_loss, z_q
+
+    def decode(self, tokens: torch.Tensor) -> torch.Tensor:
+        """
+        Decode token sequence back to chart features.
+
+        Args:
+            tokens: [B, T] — discrete token IDs.
+
+        Returns:
+            x_recon: [B, T, input_dim] — reconstructed note features.
+        """
+        z_q = self.vq.decode_indices(tokens)
+        x_recon = self.decoder(z_q)
+        return x_recon
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Full forward pass: encode → quantize → decode.
+
+        Args:
+            x: [B, T, input_dim]
+
+        Returns:
+            x_recon:  [B, T, input_dim] reconstructed features.
+            tokens:   [B, T] discrete token IDs.
+            vq_loss:  scalar loss.
+        """
+        tokens, vq_loss, z_q = self.encode(x)
+        x_recon = self.decoder(z_q)
+        return x_recon, tokens, vq_loss
+
+    @property
+    def vocab_size(self) -> int:
+        return self.codebook_size

Tokenizer/MaiTrackTokenizer.py

@@ -0,0 +1,422 @@
+"""
+MaiTrackTokenizer — Audio tokenizer using pre-trained EnCodec (Meta).
+
+Wraps EnCodec 24kHz model to convert maimai track.mp3 audio into
+discrete token sequences for transformer training.
+
+Key design decisions:
+  - Uses EnCodec 24kHz (pre-trained on speech+music, 8-layer RVQ, 1024 bins)
+  - Default: 2 codebook layers (~15816 tokens/2min song)
+  - Stride: 320 samples @ 24kHz = 75Hz = 13.3ms per token
+  - BPM is NOT encoded (computed separately by external program)
+
+Usage:
+    from Tokenizer.MaiTrackTokenizer import MaiTrackTokenizer
+
+    tok = MaiTrackTokenizer()
+    tokens = tok.encode("datasets/10/track.mp3")         # → list[int]
+    tokens_2l = tok.encode("datasets/10/track.mp3", n_layers=2)
+    audio = tok.decode(tokens)                            # → torch.Tensor
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Optional, Union
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from encodec import EncodecModel
+from encodec.utils import convert_audio
+
+logger = logging.getLogger(__name__)
+
+# ═══════════════════════════════════════════════════════════════════════
+# Constants
+# ═══════════════════════════════════════════════════════════════════════
+
+# EnCodec 24kHz: 8 codebooks, 1024 bins each, stride=320 samples
+# Frame rate = 24000 / 320 = 75 Hz (= 13.33 ms per token)
+ENC_SAMPLE_RATE = 24000
+ENC_STRIDE = 320          # samples per frame
+ENC_FRAME_RATE = ENC_SAMPLE_RATE / ENC_STRIDE  # 75 Hz
+ENC_CODEBOOK_SIZE = 1024
+ENC_NUM_CODEBOOKS = 8
+ENC_BANDWIDTH = 6.0       # kbps target
+
+# Default: use only first 2 codebook layers for efficiency
+DEFAULT_N_LAYERS = 2
+
+# Special tokens (aligned with chart tokenizer)
+PAD = 0
+BOS = 1
+EOS = 2
+
+# Token value offset per layer (to avoid overlap)
+# Layer 0: codes 0..1023 → tokens 3..1026
+# Layer 1: codes 0..1023 → tokens 1027..2050
+# etc.
+TOKEN_OFFSET_BASE = 3  # after PAD, BOS, EOS
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Tokenizer class
+# ═══════════════════════════════════════════════════════════════════════
+
+class MaiTrackTokenizer:
+    """
+    Pre-trained audio tokenizer using Meta EnCodec.
+
+    Converts audio waveforms to/from discrete token sequences.
+    Multi-layer tokens are interleaved: [L0_t0, L1_t0, L0_t1, L1_t1, ...]
+
+    Attributes:
+        sample_rate: 24000 Hz
+        frame_rate: 75 Hz (13.3ms per token)
+        n_layers: Number of codebook layers used (default 2)
+        vocab_size: Total vocabulary size (layers × 1024 + special tokens)
+    """
+
+    def __init__(self, n_layers: int = DEFAULT_N_LAYERS, device: str = "cpu"):
+        """
+        Args:
+            n_layers: Number of EnCodec codebook layers to use (1-8).
+                      More layers = better audio quality, more tokens.
+                      1-2 layers typically sufficient for rhythm game features.
+            device: Device to run the model on ("cpu" or "cuda").
+        """
+        self.n_layers = n_layers
+        self.device = device
+
+        # Load pre-trained model
+        self._model = EncodecModel.encodec_model_24khz()
+        self._model.set_target_bandwidth(ENC_BANDWIDTH)
+        self._model.eval()
+        self._model.to(device)
+
+        # Vocabulary: specials + n_layers * 1024 codes
+        self.vocab_size = TOKEN_OFFSET_BASE + n_layers * ENC_CODEBOOK_SIZE
+        self.pad_token_id = PAD
+        self.bos_token_id = BOS
+        self.eos_token_id = EOS
+
+    @property
+    def sample_rate(self) -> int:
+        return ENC_SAMPLE_RATE
+
+    @property
+    def frame_rate(self) -> float:
+        return ENC_FRAME_RATE
+
+    # ── Load audio ──────────────────────────────────────────────────
+
+    def load_audio(self, path: Union[str, Path]) -> torch.Tensor:
+        """
+        Load an audio file and convert to 24kHz mono tensor.
+
+        Args:
+            path: Path to audio file (mp3, wav, flac, etc.).
+
+        Returns:
+            Tensor [1, samples] at 24kHz mono.
+        """
+        data, sr = sf.read(str(path), dtype="float32")
+
+        # Convert to mono
+        if data.ndim > 1:
+            data = data.mean(axis=1)
+
+        wav = torch.from_numpy(data.copy()).unsqueeze(0)  # [1, T]
+
+        # Resample to 24kHz if needed
+        if sr != ENC_SAMPLE_RATE:
+            wav = convert_audio(wav, sr, ENC_SAMPLE_RATE, 1)
+
+        return wav
+
+    def load_audio_batch(self, paths: list[str],
+                         max_duration: Optional[float] = None) -> tuple[torch.Tensor, list[int]]:
+        """
+        Load a batch of audio files, padding to same length.
+
+        Args:
+            paths: List of audio file paths.
+            max_duration: Truncate to max_duration seconds (None = no truncation).
+
+        Returns:
+            (wavs [B, 1, max_samples], lengths [B])
+        """
+        wavs = []
+        lengths = []
+        for p in paths:
+            wav = self.load_audio(p)
+            if max_duration is not None:
+                max_samples = int(max_duration * ENC_SAMPLE_RATE)
+                wav = wav[:, :max_samples]
+            lengths.append(wav.shape[1])
+            wavs.append(wav)
+
+        # Pad to max length
+        max_len = max(lengths)
+        padded = torch.zeros(len(wavs), 1, max_len)
+        for i, w in enumerate(wavs):
+            padded[i, :, :w.shape[1]] = w
+
+        return padded, lengths
+
+    # ── Encode ──────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def encode(self, audio: Union[str, Path, torch.Tensor, np.ndarray],
+               n_layers: Optional[int] = None,
+               add_bos: bool = True,
+               add_eos: bool = True,
+               interleave: bool = True) -> list[int]:
+        """
+        Encode audio into a discrete token sequence.
+
+        Args:
+            audio: Path to audio file, or waveform tensor [1, T] / numpy [T].
+            n_layers: Override number of codebook layers (default: self.n_layers).
+            add_bos: Prepend BOS token.
+            add_eos: Append EOS token.
+            interleave: If True, interleave layers: [L0, L1, L0, L1, ...].
+                        If False, concatenate: [L0_all..., L1_all...].
+
+        Returns:
+            List of integer token IDs.
+        """
+        n_layers = n_layers or self.n_layers
+
+        # Load if path
+        if isinstance(audio, (str, Path)):
+            wav = self.load_audio(audio).to(self.device)
+        elif isinstance(audio, np.ndarray):
+            wav = torch.from_numpy(audio.astype("float32")).unsqueeze(0).to(self.device)
+        else:
+            wav = audio.to(self.device)
+
+        # Ensure correct shape
+        if wav.dim() == 1:
+            wav = wav.unsqueeze(0).unsqueeze(0)
+        elif wav.dim() == 2:
+            wav = wav.unsqueeze(1)
+
+        # Resample if needed
+        if wav.shape[1] != 1:
+            wav = wav.mean(dim=1, keepdim=True)
+
+        # Encode with EnCodec. Long files are returned as multiple chunks;
+        # concatenate them so token length reflects the full track duration.
+        encoded = self._model.encode(wav)
+        codes = torch.cat([frame_codes for frame_codes, _ in encoded], dim=-1)  # [B, n_q, T]
+        codes = codes[0, :n_layers, :]  # [n_layers, T]
+
+        # Convert to flat token list
+        tokens: list[int] = []
+        if add_bos:
+            tokens.append(BOS)
+
+        T = codes.shape[1]
+        if interleave:
+            # Interleave: L0_t0, L1_t0, L0_t1, L1_t1, ...
+            for t in range(T):
+                for layer in range(n_layers):
+                    code = codes[layer, t].item()
+                    token = self._code_to_token(code, layer)
+                    tokens.append(token)
+        else:
+            # Concatenate per layer
+            for layer in range(n_layers):
+                for t in range(T):
+                    code = codes[layer, t].item()
+                    token = self._code_to_token(code, layer)
+                    tokens.append(token)
+
+        if add_eos:
+            tokens.append(EOS)
+
+        return tokens
+
+    @torch.no_grad()
+    def encode_batch(self, audios: list[Union[str, Path, torch.Tensor]],
+                     n_layers: Optional[int] = None,
+                     max_duration: Optional[float] = None,
+                     pad_to: Optional[int] = None,
+                     return_tensors: bool = False):
+        """
+        Encode a batch of audio files.
+
+        Args:
+            audios: List of paths or tensors.
+            n_layers: Number of codebook layers.
+            max_duration: Truncate audio to this many seconds.
+            pad_to: Pad token sequences to this length.
+            return_tensors: Return torch tensors instead of lists.
+
+        Returns:
+            If return_tensors=False: (list[list[int]], list[int])
+            If return_tensors=True: (Tensor[B, max_len], Tensor[B])
+        """
+        n_layers = n_layers or self.n_layers
+
+        token_seqs = []
+        for audio in audios:
+            tokens = self.encode(audio, n_layers=n_layers, interleave=True)
+            if pad_to is not None and len(tokens) > pad_to:
+                tokens = tokens[:pad_to]
+            token_seqs.append(tokens)
+
+        lengths = [len(s) for s in token_seqs]
+        max_len = max(lengths) if pad_to is None else max(pad_to, max(lengths))
+
+        padded = []
+        for seq in token_seqs:
+            if len(seq) < max_len:
+                seq = seq + [PAD] * (max_len - len(seq))
+            padded.append(seq[:max_len])
+
+        if return_tensors:
+            return (torch.tensor(padded, dtype=torch.long),
+                    torch.tensor(lengths, dtype=torch.long))
+
+        return padded, lengths
+
+    # ── Decode ──────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def decode(self, tokens: list[int],
+               n_layers: Optional[int] = None,
+               interleave: bool = True) -> torch.Tensor:
+        """
+        Decode a token sequence back to audio waveform.
+
+        Args:
+            tokens: Token ID list.
+            n_layers: Number of codebook layers used (must match encoding).
+            interleave: Whether tokens are interleaved (must match encoding).
+
+        Returns:
+            Waveform tensor [1, samples] at 24kHz.
+        """
+        n_layers = n_layers or self.n_layers
+
+        # Strip BOS/EOS
+        if tokens and tokens[0] == BOS:
+            tokens = tokens[1:]
+        if tokens and tokens[-1] == EOS:
+            tokens = tokens[:-1]
+
+        total_tokens = len(tokens)
+        if interleave:
+            T = total_tokens // n_layers
+        else:
+            T = total_tokens // n_layers
+
+        if T == 0:
+            logger.warning("Token sequence too short for decoding")
+            return torch.zeros(1, ENC_STRIDE)
+
+        # Convert tokens back to EnCodec codes
+        # Full 8 layers: fill unused layers with zeros
+        codes = torch.zeros(ENC_NUM_CODEBOOKS, T, dtype=torch.long, device=self.device)
+
+        if interleave:
+            for t in range(T):
+                for layer in range(n_layers):
+                    idx = t * n_layers + layer
+                    if idx < total_tokens:
+                        codes[layer, t] = self._token_to_code(tokens[idx], layer)
+        else:
+            for layer in range(n_layers):
+                for t in range(T):
+                    idx = layer * T + t
+                    if idx < total_tokens:
+                        codes[layer, t] = self._token_to_code(tokens[idx], layer)
+
+        # Decode with EnCodec
+        codes = codes.unsqueeze(0)  # [1, 8, T]
+        decoded = self._model.decode([(codes, None)])
+        return decoded.squeeze(0)  # [1, samples]
+
+    # ── Token ↔ Code conversion ─────────────────────────────────────
+
+    def _code_to_token(self, code: int, layer: int) -> int:
+        """Convert EnCodec code (0..1023) to global token ID."""
+        return TOKEN_OFFSET_BASE + layer * ENC_CODEBOOK_SIZE + code
+
+    def _token_to_code(self, token: int, layer: int) -> int:
+        """Convert global token ID back to EnCodec code (0..1023)."""
+        offset = TOKEN_OFFSET_BASE + layer * ENC_CODEBOOK_SIZE
+        code = token - offset
+        return max(0, min(code, ENC_CODEBOOK_SIZE - 1))
+
+    # ── Debug ───────────────────────────────────────────────────────
+
+    def tokens_to_str(self, tokens: list[int], max_show: int = 30) -> str:
+        """Pretty-print token sequence (truncated)."""
+        parts = []
+        for t in tokens[:max_show]:
+            if t == PAD:
+                parts.append("[PAD]")
+            elif t == BOS:
+                parts.append("[BOS]")
+            elif t == EOS:
+                parts.append("[EOS]")
+            else:
+                # Determine layer and code
+                code = t - TOKEN_OFFSET_BASE
+                layer = code // ENC_CODEBOOK_SIZE
+                c = code % ENC_CODEBOOK_SIZE
+                parts.append(f"L{layer}:{c}")
+        if len(tokens) > max_show:
+            parts.append(f"... ({len(tokens) - max_show} more)")
+        return " ".join(parts)
+
+    # ── Info ────────────────────────────────────────────────────────
+
+    def token_count_estimate(self, duration_seconds: float,
+                             n_layers: Optional[int] = None) -> int:
+        """
+        Estimate the number of tokens for a given audio duration.
+
+        Args:
+            duration_seconds: Audio duration in seconds.
+            n_layers: Number of codebook layers.
+
+        Returns:
+            Estimated token count (including BOS/EOS).
+        """
+        n_layers = n_layers or self.n_layers
+        frames = int(duration_seconds * ENC_FRAME_RATE)
+        return 2 + frames * n_layers  # BOS + frames*layers + EOS
+
+    def __repr__(self) -> str:
+        return (f"MaiTrackTokenizer(n_layers={self.n_layers}, "
+                f"sr={ENC_SAMPLE_RATE}Hz, "
+                f"frame_rate={ENC_FRAME_RATE:.0f}Hz, "
+                f"vocab_size={self.vocab_size}, "
+                f"device={self.device})")
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Quick test
+# ═══════════════════════════════════════════════════════════════════════
+
+if __name__ == "__main__":
+    import sys
+
+    path = sys.argv[1] if len(sys.argv) > 1 else "datasets/10/track.mp3"
+
+    tok = MaiTrackTokenizer(n_layers=2)
+    print(tok)
+    print(f"Vocab size: {tok.vocab_size}")
+
+    tokens = tok.encode(path)
+    print(f"Tokens: {len(tokens)} ({tok.tokens_to_str(tokens, 30)})")
+
+    audio = tok.decode(tokens)
+    print(f"Decoded audio: {audio.shape}, {audio.shape[1]/ENC_SAMPLE_RATE:.1f}s")

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:886781211e93a4d3c30dceb456ae21c1f456f50e2f086f484bd53505e879c3ec
+size 1599244629

constrained_decode.py

@@ -0,0 +1,389 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from grammar import ChartGrammarState, GrammarMode
+from model import DIV_MAP, DUR_TOKEN, is_timeline_token
+from grammar import NOTE_TOKENS
+from tokenizer import (
+    BOS, EOS, PAD, RST, DUR, HLD_TO_ID, ID_TO_DIV, SIM_BEG, SIM_END, SLD_BEG,
+    SLD_END_TOKEN, SIM_COUNT_2, DUR_NUM_TO_ID, ID_TO_DUR_NUM, DUR_DEN_TO_ID,
+    ID_TO_DUR_DEN, CONFIG_TO_ID, ID_TO_CONFIG,
+)
+
+
+VALID_DUR_NUMS = list(DUR_NUM_TO_ID.keys())
+VALID_DUR_DENS = list(DUR_DEN_TO_ID.keys())
+MAX_DURATION_BEATS = 4.0
+
+
+def snap_duration(beat: int, subdiv: int, max_beats: float,
+                  max_duration_beats: float = MAX_DURATION_BEATS) -> tuple[int, int]:
+    """Snap/clamp duration params to common values that fit before track end."""
+    limit = min(max_beats, max_duration_beats)
+    if limit <= 0:
+        return 1, 64
+    candidates = [
+        (b, d) for b in VALID_DUR_NUMS for d in VALID_DUR_DENS
+        if b / d <= limit + 1e-6
+    ]
+    if not candidates:
+        return 1, 64
+    raw = max(float(beat), 1.0) / max(float(subdiv), 1.0)
+    return min(candidates, key=lambda x: (abs((x[0] / x[1]) - raw), x[1], x[0]))
+
+
+def _config_duration_beats(spec: tuple) -> float:
+    if not spec:
+        return 0.0
+    kind = spec[0]
+    if kind in ("hld", "sld") and len(spec) >= 4:
+        return max(float(spec[-2]), 1.0) / max(float(spec[-1]), 1.0)
+    if kind == "pair" and len(spec) >= 7 and ("hld" in (spec[1], spec[3])):
+        return max(float(spec[-2]), 1.0) / max(float(spec[-1]), 1.0)
+    if kind == "multi" and len(spec) >= 5 and spec[1] == "hld":
+        return max(float(spec[-2]), 1.0) / max(float(spec[-1]), 1.0)
+    return 0.0
+
+
+def _replace_config_duration_token(token_id: int, max_beats: float,
+                                   max_duration_beats: float = MAX_DURATION_BEATS) -> int:
+    spec = ID_TO_CONFIG.get(token_id)
+    if spec is None or _config_duration_beats(spec) <= 0:
+        return token_id
+    n, d = snap_duration(int(spec[-2]), int(spec[-1]), max_beats, max_duration_beats)
+    new_spec = tuple(list(spec[:-2]) + [n, d])
+    return CONFIG_TO_ID.get(new_spec, RST)
+
+
+@dataclass
+class DecodeState:
+    """Stateful validator for candidate next-token decoding."""
+
+    bpm: float
+    total_beats: float
+    start_offset_beats: float
+    min_division: int = 4
+    max_sim: int = 2
+    max_duration_beats: float = MAX_DURATION_BEATS
+
+    def __post_init__(self) -> None:
+        self.grammar = ChartGrammarState(min_division=self.min_division, max_sim=self.max_sim)
+        self.current_beat = float(self.start_offset_beats)
+        self.div_value = float(self.min_division)
+        self.skip_beat = 0
+        self.slide_group_active = False
+        self.pending_dur_num: int | None = None
+        self.pending_duration_start_beat: float | None = None
+        self.sim_start_beat: float | None = None
+        self.slide_start_beat: float | None = None
+        self._normal_mask_cache: dict[tuple[torch.device, int], torch.Tensor] = {}
+        self._config_duration_cache: dict[tuple[torch.device, int], torch.Tensor] = {}
+
+    def remaining_beats(self) -> float:
+        return max(0.0, self.total_beats - self.current_beat)
+
+    def remaining_duration_beats(self) -> float:
+        start = self.pending_duration_start_beat
+        if start is None:
+            start = self.current_beat
+        return max(0.0, min(self.total_beats - start, self.max_duration_beats))
+
+    def allowed_tokens(self) -> set[int]:
+        allowed = set(self.grammar.allowed_tokens())
+        remaining = self.remaining_duration_beats()
+
+        if self.grammar.mode is GrammarMode.DUR_NUM:
+            nums = {DUR_NUM_TO_ID[b] for b in VALID_DUR_NUMS if b / max(VALID_DUR_DENS) <= remaining + 1e-6}
+            return allowed & nums if nums else {DUR_NUM_TO_ID[1]}
+
+        if self.grammar.mode is GrammarMode.DUR_DEN:
+            beat_num = self.pending_dur_num or 1
+            dens = {DUR_DEN_TO_ID[d] for d in VALID_DUR_DENS if beat_num / d <= remaining + 1e-6}
+            return allowed & dens if dens else {DUR_DEN_TO_ID[64]}
+
+        filtered: set[int] = set()
+        for tok in allowed:
+            if self.is_candidate_safe(tok):
+                filtered.add(tok)
+        return filtered
+
+    def can_start_sim(self) -> bool:
+        return self.grammar.mode is GrammarMode.NORMAL and self.max_sim >= 2
+
+    def is_candidate_safe(self, tok: int) -> bool:
+        if self.grammar.mode is not GrammarMode.NORMAL:
+            return True
+        if tok in (PAD, BOS, EOS):
+            return False
+        if tok == SIM_BEG and self.max_sim < 2:
+            return False
+        if tok in ID_TO_DIV:
+            next_div = float(ID_TO_DIV[tok])
+            return self.remaining_beats() >= 4.0 / max(next_div, 1.0) - 1e-6
+        if tok in HLD_TO_ID.values() or tok == SLD_BEG:
+            if self.remaining_beats() < 1.0 / max(VALID_DUR_DENS):
+                return False
+        if self._would_advance_time(tok):
+            step = 4.0 / max(self.div_value, 1.0)
+            return self.current_beat + step <= self.total_beats + 1e-6
+        return True
+
+    def apply_logits_mask(self, logits: torch.Tensor) -> torch.Tensor:
+        masked = torch.full_like(logits, float("-inf"))
+        vocab_size = int(logits.numel())
+
+        if self.grammar.mode is GrammarMode.NORMAL:
+            mask = self._normal_mask(logits.device, vocab_size).clone()
+
+            step = 4.0 / max(self.div_value, 1.0)
+            if self.current_beat + step > self.total_beats + 1e-6:
+                timeline_ids = [RST, SIM_BEG, SLD_BEG, *NOTE_TOKENS, *ID_TO_CONFIG.keys()]
+                timeline_ids = [t for t in timeline_ids if 0 <= int(t) < vocab_size]
+                if timeline_ids:
+                    mask[torch.tensor(timeline_ids, dtype=torch.long, device=logits.device)] = False
+
+            remaining = self.remaining_duration_beats()
+            if remaining < 1.0 / max(VALID_DUR_DENS):
+                risky_ids = [SLD_BEG, *HLD_TO_ID.values()]
+                risky_ids = [t for t in risky_ids if 0 <= int(t) < vocab_size]
+                if risky_ids:
+                    mask[torch.tensor(risky_ids, dtype=torch.long, device=logits.device)] = False
+
+            cfg_dur = self._config_duration_tensor(logits.device, vocab_size)
+            mask &= cfg_dur <= remaining + 1e-6
+
+            for div_tok, div_value in ID_TO_DIV.items():
+                if 0 <= div_tok < vocab_size:
+                    mask[div_tok] = self.remaining_beats() >= 4.0 / max(float(div_value), 1.0) - 1e-6
+
+            masked[mask] = logits[mask]
+            return masked
+
+        allowed = self.allowed_tokens()
+        if allowed:
+            idx = torch.tensor([t for t in sorted(allowed) if 0 <= t < vocab_size],
+                               dtype=torch.long, device=logits.device)
+            if idx.numel() > 0:
+                masked[idx] = logits[idx]
+        return masked
+
+    def _normal_mask(self, device: torch.device, vocab_size: int) -> torch.Tensor:
+        key = (device, vocab_size)
+        cached = self._normal_mask_cache.get(key)
+        if cached is not None:
+            return cached
+
+        mask = torch.zeros(vocab_size, dtype=torch.bool, device=device)
+        ids = {RST, SLD_BEG} | set(NOTE_TOKENS) | set(ID_TO_CONFIG.keys()) | set(self.grammar.allowed_div_tokens)
+        if self.max_sim >= 2:
+            ids.add(SIM_BEG)
+        ids.discard(PAD)
+        ids.discard(BOS)
+        ids.discard(EOS)
+        valid = [int(t) for t in ids if 0 <= int(t) < vocab_size]
+        if valid:
+            mask[torch.tensor(valid, dtype=torch.long, device=device)] = True
+        self._normal_mask_cache[key] = mask
+        return mask
+
+    def _config_duration_tensor(self, device: torch.device, vocab_size: int) -> torch.Tensor:
+        key = (device, vocab_size)
+        cached = self._config_duration_cache.get(key)
+        if cached is not None:
+            return cached
+
+        dur = torch.zeros(vocab_size, dtype=torch.float32, device=device)
+        for token_id, spec in ID_TO_CONFIG.items():
+            if 0 <= token_id < vocab_size:
+                d = _config_duration_beats(spec)
+                if d > 0:
+                    dur[token_id] = float(d)
+        self._config_duration_cache[key] = dur
+        return dur
+
+    def step(self, tok: int) -> None:
+        mode_before = self.grammar.mode
+
+        if mode_before is GrammarMode.NORMAL:
+            if tok in HLD_TO_ID.values():
+                self.pending_duration_start_beat = self.current_beat
+            elif tok == SLD_BEG:
+                self.slide_start_beat = self.current_beat
+            elif tok == SIM_BEG:
+                self.sim_start_beat = self.current_beat
+
+        if mode_before is GrammarMode.SLIDE_BODY and tok == SLD_END_TOKEN:
+            self.pending_duration_start_beat = self.slide_start_beat
+            self.slide_start_beat = None
+
+        if mode_before is GrammarMode.SIM_BODY and tok == SIM_END and self.grammar.sim_has_hold:
+            self.pending_duration_start_beat = self.sim_start_beat
+            self.sim_start_beat = None
+
+        if mode_before is GrammarMode.DUR_NUM:
+            self.pending_dur_num = ID_TO_DUR_NUM.get(int(tok), int(tok))
+        elif mode_before is GrammarMode.DUR_DEN and self.pending_dur_num is not None:
+            den_value = ID_TO_DUR_DEN.get(int(tok), max(int(tok), 1))
+            beat, den = snap_duration(self.pending_dur_num, den_value,
+                                      self.remaining_duration_beats(),
+                                      self.max_duration_beats)
+            self.pending_dur_num = None
+            self.pending_duration_start_beat = None
+
+        in_sim_body = self.grammar.mode in (GrammarMode.SIM_COUNT, GrammarMode.SIM_BODY)
+        in_slide_body = self.slide_group_active and tok != SLD_BEG
+        if self.skip_beat > 0:
+            self.skip_beat -= 1
+        elif tok == DUR_TOKEN:
+            self.skip_beat = 2
+        elif tok in ID_TO_DIV:
+            self.div_value = float(ID_TO_DIV[tok])
+        elif is_timeline_token(torch.tensor(tok)) and not in_sim_body and not in_slide_body:
+            self.current_beat += 4.0 / max(self.div_value, 1.0)
+
+        if tok == SLD_BEG:
+            self.slide_group_active = True
+        elif tok == SLD_END_TOKEN:
+            self.slide_group_active = False
+
+        self.grammar.step(tok)
+
+    def _would_advance_time(self, tok: int) -> bool:
+        if self.skip_beat > 0 or tok == DUR_TOKEN:
+            return False
+        in_sim_body = self.grammar.mode in (GrammarMode.SIM_COUNT, GrammarMode.SIM_BODY)
+        in_slide_body = self.slide_group_active and tok != SLD_BEG
+        return bool(is_timeline_token(torch.tensor(tok)) and not in_sim_body and not in_slide_body)
+
+
+def clamp_duration_tokens(tokens: list[int], total_beats: float,
+                          start_offset_beats: float, min_division: int = 4) -> list[int]:
+    """Clamp generated hold/slide durations so notes cannot extend past track end."""
+    result = list(tokens)
+    state = DecodeState(bpm=150.0, total_beats=total_beats,
+                        start_offset_beats=start_offset_beats,
+                        min_division=min_division)
+    i = 0
+    while i < len(result):
+        tok = result[i]
+        if tok in ID_TO_CONFIG:
+            result[i] = _replace_config_duration_token(
+                tok,
+                state.remaining_duration_beats(),
+                state.max_duration_beats,
+            )
+            state.step(result[i])
+            i += 1
+            continue
+        if tok == DUR and i + 2 < len(result):
+            result[i + 1], result[i + 2] = snap_duration(
+                ID_TO_DUR_NUM.get(result[i + 1], result[i + 1]),
+                ID_TO_DUR_DEN.get(result[i + 2], result[i + 2]),
+                state.remaining_duration_beats(),
+                state.max_duration_beats
+            )
+            result[i + 1] = DUR_NUM_TO_ID[result[i + 1]]
+            result[i + 2] = DUR_DEN_TO_ID[result[i + 2]]
+            state.step(tok)
+            state.step(result[i + 1])
+            state.step(result[i + 2])
+            i += 3
+            continue
+        state.step(tok)
+        i += 1
+    return result
+
+
+def duration_report(tokens: list[int], total_beats: float, start_offset_beats: float,
+                    min_division: int = 4) -> dict[str, float | int]:
+    state = DecodeState(bpm=150.0, total_beats=total_beats,
+                        start_offset_beats=start_offset_beats,
+                        min_division=min_division)
+    n_dur = 0
+    max_dur = 0.0
+    overrun = 0
+    i = 0
+    while i < len(tokens):
+        tok = tokens[i]
+        if tok in ID_TO_CONFIG:
+            dur = _config_duration_beats(ID_TO_CONFIG[tok])
+            if dur > 0:
+                n_dur += 1
+                max_dur = max(max_dur, dur)
+                if state.current_beat + dur > total_beats + 1e-6:
+                    overrun += 1
+        if tok == DUR and i + 2 < len(tokens):
+            num = ID_TO_DUR_NUM.get(tokens[i + 1], tokens[i + 1])
+            den = ID_TO_DUR_DEN.get(tokens[i + 2], tokens[i + 2])
+            dur = max(float(num), 1.0) / max(float(den), 1.0)
+            n_dur += 1
+            max_dur = max(max_dur, dur)
+            start = state.pending_duration_start_beat
+            if start is None:
+                start = state.current_beat
+            if start + dur > total_beats + 1e-6:
+                overrun += 1
+        state.step(tok)
+        i += 1
+    return {"durations": n_dur, "max_beats": max_dur, "overrun": overrun}
+
+
+def simultaneous_report(tokens: list[int], max_sim: int = 2) -> dict[str, int]:
+    n_sim = 0
+    max_count = 0
+    over_limit = 0
+    first_bad = -1
+    i = 0
+    while i < len(tokens):
+        tok = tokens[i]
+        if tok == SIM_BEG and i + 1 < len(tokens):
+            count = 2 if tokens[i + 1] == SIM_COUNT_2 else int(tokens[i + 1])
+            n_sim += 1
+            max_count = max(max_count, count)
+            if count > max_sim:
+                over_limit += 1
+                if first_bad < 0:
+                    first_bad = i
+            i += 2
+            continue
+        i += 1
+    return {"groups": n_sim, "max_count": max_count, "over_limit": over_limit, "first_bad": first_bad}
+
+
+def sanitize_sim_tokens(tokens: list[int], max_sim: int = 2) -> list[int]:
+    """Remove malformed or over-wide SIM groups before decode/export."""
+    result: list[int] = []
+    i = 0
+    while i < len(tokens):
+        tok = tokens[i]
+        if tok != SIM_BEG:
+            result.append(tok)
+            i += 1
+            continue
+
+        if i + 1 >= len(tokens):
+            i += 1
+            continue
+
+        declared = int(tokens[i + 1])
+        j = i + 2
+        body: list[int] = []
+        while j < len(tokens) and tokens[j] not in (SIM_END, EOS):
+            if tokens[j] in NOTE_TOKENS:
+                body.append(tokens[j])
+            j += 1
+
+        declared_count = 2 if declared == SIM_COUNT_2 else declared
+        if declared_count == 2 and len(body) == 2 and max_sim >= 2:
+            result.extend([SIM_BEG, SIM_COUNT_2, body[0], body[1], SIM_END])
+        elif len(body) >= 2 and max_sim >= 2:
+            result.extend([SIM_BEG, SIM_COUNT_2, body[0], body[1], SIM_END])
+            result.extend(body[2:])
+        else:
+            result.extend(body[:1])
+
+        i = j + 1 if j < len(tokens) and tokens[j] == SIM_END else j
+    return result

grammar.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum, auto
+
+import torch
+from tokenizer import (
+    BOS, EOS, HLD_ON, MASK, META_BPM, META_END, PAD, RST, SEP, SLD_MID, SLD_ON,
+    BRK_TO_ID, DUR, HLD_TO_ID, ID_TO_DIV, SIM_BEG, SIM_END, SLD_BEG,
+    SLD_END_TOKEN, SLD_TO_ID, TAP_TO_ID, TCH_TO_ID,
+    SIM_COUNT_2, DUR_NUM_TO_ID, DUR_DEN_TO_ID,
+    CONFIG_BASE,
+)
+
+
+class GrammarMode(Enum):
+    NORMAL = auto()
+    SIM_COUNT = auto()
+    SIM_BODY = auto()
+    NEED_DUR = auto()
+    DUR_NUM = auto()
+    DUR_DEN = auto()
+    SLIDE_BODY = auto()
+
+
+NOTE_TOKENS = (
+    set(TAP_TO_ID.values())
+    | set(BRK_TO_ID.values())
+    | set(HLD_TO_ID.values())
+    | set(SLD_TO_ID.values())
+    | set(TCH_TO_ID.values())
+)
+HOLD_TOKENS = set(HLD_TO_ID.values())
+SLIDE_POINT_TOKENS = set(SLD_TO_ID.values())
+META_TOKENS = set(range(META_BPM, META_END + 1))
+ENGINE_ONLY_TOKENS = {PAD, BOS, EOS, SEP, MASK, HLD_ON, SLD_ON, SLD_MID} | META_TOKENS
+# Duration parameters are raw integer tokens, not normal vocabulary tokens.
+# Keep them in the same compact domain produced by real charts; allowing the
+# whole uint8 range lets generation sample special/control-looking ids and
+# absurdly long holds/slides.
+VALID_DUR_NUMS = set(DUR_NUM_TO_ID.values())
+VALID_DUR_DENS = set(DUR_DEN_TO_ID.values())
+
+
+def config_tokens() -> set[int]:
+    # Config vocab can be extended after importing this module when a checkpoint
+    # loads its learned config_vocab. Keep this dynamic instead of freezing at import.
+    import tokenizer as chart_tokenizer
+    return set(chart_tokenizer.ID_TO_CONFIG.keys())
+
+
+@dataclass
+class ChartGrammarState:
+    """Hard token grammar for chart generation.
+
+    This keeps the sequence parseable. It intentionally does not enforce chart
+    quality or muri rules; those belong in a later event/window validator.
+    """
+
+    min_division: int = 1
+    max_sim: int = 2
+    mode: GrammarMode = GrammarMode.NORMAL
+    sim_target: int = 0
+    sim_seen: int = 0
+    sim_has_hold: bool = False
+    slide_points: int = 0
+    last_event_token: int | None = None
+    can_take_duration: bool = False
+    active_slide: bool = False
+    allowed_div_tokens: set[int] = field(default_factory=set)
+
+    def __post_init__(self) -> None:
+        self.allowed_div_tokens = {
+            div_id for div_id, div_value in ID_TO_DIV.items()
+            if div_value >= self.min_division
+        }
+
+    def allowed_tokens(self) -> set[int]:
+        if self.mode is GrammarMode.SIM_COUNT:
+            return {SIM_COUNT_2} if self.max_sim >= 2 else set()
+
+        if self.mode is GrammarMode.SIM_BODY:
+            if self.sim_seen >= self.sim_target:
+                return {SIM_END}
+            return NOTE_TOKENS - SLIDE_POINT_TOKENS
+
+        if self.mode is GrammarMode.NEED_DUR:
+            return {DUR}
+
+        if self.mode is GrammarMode.DUR_NUM:
+            return VALID_DUR_NUMS
+
+        if self.mode is GrammarMode.DUR_DEN:
+            return VALID_DUR_DENS
+
+        if self.mode is GrammarMode.SLIDE_BODY:
+            allowed = set(SLIDE_POINT_TOKENS)
+            if self.slide_points >= 2:
+                allowed.add(SLD_END_TOKEN)
+            return allowed
+
+        allowed = {RST, SIM_BEG, SLD_BEG} | NOTE_TOKENS | config_tokens() | self.allowed_div_tokens
+        return allowed - ENGINE_ONLY_TOKENS
+
+    def apply_logits_mask(self, logits: torch.Tensor) -> torch.Tensor:
+        allowed = self.allowed_tokens()
+        masked = torch.full_like(logits, float("-inf"))
+        idx = torch.tensor(sorted(allowed), dtype=torch.long, device=logits.device)
+        masked[idx] = logits[idx]
+        return masked
+
+    def step(self, token: int) -> None:
+        if self.mode is GrammarMode.SIM_COUNT:
+            self.sim_target = 2 if token == SIM_COUNT_2 else min(max(int(token), 1), self.max_sim)
+            self.sim_seen = 0
+            self.sim_has_hold = False
+            self.mode = GrammarMode.SIM_BODY
+            return
+
+        if self.mode is GrammarMode.SIM_BODY:
+            if token == SIM_END:
+                self.last_event_token = SIM_END
+                if self.sim_has_hold:
+                    self.mode = GrammarMode.NEED_DUR
+                    self.can_take_duration = True
+                else:
+                    self.mode = GrammarMode.NORMAL
+                    self.can_take_duration = False
+            else:
+                if token in HOLD_TOKENS:
+                    self.sim_has_hold = True
+                self.sim_seen += 1
+            return
+
+        if self.mode is GrammarMode.NEED_DUR:
+            if token == DUR:
+                self.mode = GrammarMode.DUR_NUM
+            return
+
+        if self.mode is GrammarMode.DUR_NUM:
+            self.mode = GrammarMode.DUR_DEN
+            return
+
+        if self.mode is GrammarMode.DUR_DEN:
+            self.mode = GrammarMode.NORMAL
+            self.last_event_token = None
+            self.can_take_duration = False
+            return
+
+        if self.mode is GrammarMode.SLIDE_BODY:
+            if token == SLD_END_TOKEN:
+                self.mode = GrammarMode.NEED_DUR
+                self.active_slide = False
+                self.last_event_token = SLD_END_TOKEN
+                self.can_take_duration = True
+            elif token in SLIDE_POINT_TOKENS:
+                self.slide_points += 1
+            return
+
+        if token == SIM_BEG:
+            self.mode = GrammarMode.SIM_COUNT
+            self.last_event_token = SIM_BEG
+            self.can_take_duration = False
+        elif token == SLD_BEG:
+            self.mode = GrammarMode.SLIDE_BODY
+            self.slide_points = 0
+            self.active_slide = True
+            self.last_event_token = SLD_BEG
+            self.can_take_duration = False
+        elif token == DUR:
+            self.mode = GrammarMode.DUR_NUM
+        elif token in ID_TO_DIV:
+            return
+        elif token == RST or token in NOTE_TOKENS or token >= CONFIG_BASE:
+            self.last_event_token = token
+            if token in HOLD_TOKENS:
+                self.mode = GrammarMode.NEED_DUR
+                self.can_take_duration = True
+            else:
+                self.can_take_duration = False

inference.py

@@ -0,0 +1,707 @@
+"""
+Inference script for MaiGenerator — generate maimai charts from audio.
+
+Usage:
+    python inference.py --checkpoint checkpoints/best.pt --audio datasets/10/track.mp3
+    python inference.py --checkpoint checkpoints/best.pt --song-id 10
+    python inference.py --checkpoint checkpoints/best.pt --audio track.mp3 --bpm 173 --diff MASTER --level 12.4
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from tqdm import tqdm
+
+from model import MaiGenerator, BOS, EOS, DUR_TOKEN, AUDIO_FRAME_RATE
+from Tokenizer.MaiTrackTokenizer import MaiTrackTokenizer
+from tokenizer import (
+    MaiChartTokenizer, tokens_to_maitext, PAD, SIM_BEG, RST,
+    TAP_TO_ID, BRK_TO_ID, HLD_TO_ID, SLD_TO_ID, TCH_TO_ID,
+)
+from tokenizer import ID_TO_DUR_NUM, ID_TO_DUR_DEN, DUR_NUM_TO_ID, DUR_DEN_TO_ID
+from tokenizer import load_config_vocab, ID_TO_CONFIG
+from mai_parser import Difficulty, parse_maidata_file
+from grammar import NOTE_TOKENS
+from constrained_decode import (
+    DecodeState, snap_duration, clamp_duration_tokens, duration_report,
+    simultaneous_report, sanitize_sim_tokens,
+)
+
+
+def sample_token(logits: torch.Tensor, temperature: float = 0.8,
+                 top_k: int = 50, top_p: float = 0.95,
+                 repetition_penalty: float = 1.0,
+                 prev_tokens: list[int] | None = None,
+                 penalty_window: int = 50) -> int:
+    """Sample a token with temperature, top-k, top-p, and repetition penalty.
+
+    Args:
+        repetition_penalty: >1.0 penalizes repeated tokens; 1.0 = off.
+        prev_tokens: Recent token history for repetition detection.
+        penalty_window: How many recent tokens to check for repeats.
+    """
+    # Repetition penalty: reduce logits for tokens that appear in recent history
+    if repetition_penalty > 1.0 and prev_tokens:
+        recent = set(prev_tokens[-penalty_window:])
+        for tok in recent:
+            logits[tok] = logits[tok] / repetition_penalty
+
+    logits = logits / max(temperature, 1e-8)
+    logits = torch.nan_to_num(logits, nan=-1e9, posinf=1e9, neginf=-1e9)
+
+    if top_k > 0:
+        k = min(top_k, logits.size(-1))
+        min_val = torch.topk(logits, k).values[-1]
+        logits[logits < min_val] = float("-inf")
+
+    if top_p < 1.0:
+        sorted_l, sorted_i = torch.sort(logits, descending=True)
+        cum = torch.cumsum(F.softmax(sorted_l, dim=-1), dim=-1)
+        mask = cum > top_p
+        mask[1:] = mask[:-1].clone()
+        mask[0] = False
+        sorted_l[mask] = float("-inf")
+        logits = sorted_l.scatter(0, sorted_i, sorted_l)
+
+    probs = F.softmax(logits, dim=-1)
+    return torch.multinomial(probs, 1).item()
+
+
+def _config_note_count(spec: tuple) -> int:
+    if not spec:
+        return 0
+    kind = spec[0]
+    if kind == "pair":
+        return 2
+    if kind == "multi":
+        return max(1, len(spec) - (4 if spec[1] == "hld" else 2))
+    if kind == "touch_multi":
+        return max(1, len(spec) - 1)
+    return 1
+
+
+def _token_note_count(tok: int) -> int:
+    if tok in NOTE_TOKENS:
+        return 1
+    spec = ID_TO_CONFIG.get(tok)
+    if spec is not None:
+        return _config_note_count(spec)
+    return 0
+
+
+def _config_has_break(spec: tuple) -> bool:
+    if not spec:
+        return False
+    if spec[0] == "brk":
+        return True
+    if spec[0] == "pair":
+        return spec[1] == "brk" or spec[3] == "brk"
+    if spec[0] == "multi":
+        return len(spec) > 1 and spec[1] == "brk"
+    return False
+
+
+def _config_type_counts(spec: tuple) -> dict[str, int]:
+    if not spec:
+        return {}
+    kind = spec[0]
+    if kind == "tap":
+        return {"tap": 1}
+    if kind == "brk":
+        return {"break": 1}
+    if kind == "hld":
+        return {"hold": 1}
+    if kind == "sld":
+        return {"slide": 1}
+    if kind in ("tch", "touch_multi"):
+        return {"touch": 1}
+    if kind == "pair":
+        counts: dict[str, int] = {}
+        for typ in (spec[1], spec[3]):
+            key = "break" if typ == "brk" else "hold" if typ == "hld" else "tap"
+            counts[key] = counts.get(key, 0) + 1
+        return counts
+    if kind == "multi":
+        typ = spec[1]
+        key = "break" if typ == "brk" else "hold" if typ == "hld" else "tap"
+        if typ == "hld":
+            n = max(1, len(spec) - 4)
+        else:
+            n = max(1, len(spec) - 2)
+        return {key: n}
+    return {}
+
+
+def _token_has_break(tok: int) -> bool:
+    if tok in BRK_TO_ID.values():
+        return True
+    spec = ID_TO_CONFIG.get(tok)
+    return spec is not None and _config_has_break(spec)
+
+
+def _break_ids(vocab_size: int, device: torch.device) -> torch.Tensor:
+    ids = set(BRK_TO_ID.values())
+    ids.update(t for t, spec in ID_TO_CONFIG.items() if _config_has_break(spec))
+    valid = [int(t) for t in ids if 0 <= int(t) < vocab_size]
+    return torch.tensor(valid, dtype=torch.long, device=device)
+
+
+def _density_note_ids(vocab_size: int, device: torch.device,
+                      include_break: bool = False) -> torch.Tensor:
+    ids = set(NOTE_TOKENS) | set(ID_TO_CONFIG.keys())
+    if not include_break:
+        ids = {t for t in ids if not _token_has_break(int(t))}
+    valid = [int(t) for t in ids if 0 <= int(t) < vocab_size]
+    return torch.tensor(valid, dtype=torch.long, device=device)
+
+
+def _type_bias_vector(vocab_size: int, device: torch.device,
+                      type_biases: dict[str, float]) -> torch.Tensor | None:
+    active = {k: float(v) for k, v in type_biases.items() if abs(float(v)) > 1e-9}
+    if not active:
+        return None
+    bias = torch.zeros(vocab_size, dtype=torch.float32, device=device)
+    token_groups = [
+        ("tap", TAP_TO_ID.values()),
+        ("hold", HLD_TO_ID.values()),
+        ("slide", SLD_TO_ID.values()),
+        ("break", BRK_TO_ID.values()),
+        ("touch", TCH_TO_ID.values()),
+    ]
+    for name, ids in token_groups:
+        value = active.get(name, 0.0)
+        if value == 0.0:
+            continue
+        valid = [int(t) for t in ids if 0 <= int(t) < vocab_size]
+        if valid:
+            bias[torch.tensor(valid, dtype=torch.long, device=device)] += value
+    if active.get("rest", 0.0) and 0 <= RST < vocab_size:
+        bias[RST] += active["rest"]
+    if active.get("sim", 0.0) and 0 <= SIM_BEG < vocab_size:
+        bias[SIM_BEG] += active["sim"]
+
+    for token_id, spec in ID_TO_CONFIG.items():
+        if not (0 <= int(token_id) < vocab_size):
+            continue
+        for name, count in _config_type_counts(spec).items():
+            value = active.get(name, 0.0)
+            if value:
+                bias[int(token_id)] += value * max(1, int(count))
+    return bias
+
+
+def _config_density_bias_vector(vocab_size: int, device: torch.device,
+                                strength: float,
+                                include_break: bool = False) -> torch.Tensor | None:
+    if abs(float(strength)) <= 1e-9:
+        return None
+    bias = torch.zeros(vocab_size, dtype=torch.float32, device=device)
+    for token_id, spec in ID_TO_CONFIG.items():
+        if not (0 <= int(token_id) < vocab_size):
+            continue
+        if not include_break and _config_has_break(spec):
+            continue
+        extra_notes = max(0, _config_note_count(spec) - 1)
+        if extra_notes > 0:
+            bias[int(token_id)] += float(strength) * extra_notes
+    return bias
+
+
+def _auto_target_notes(level_value: float, duration_sec: float) -> int:
+    # Conservative maimai-ish density prior. It is only used when density bias is enabled.
+    notes_per_sec = float(np.interp(
+        level_value,
+        [10.0, 12.0, 13.0, 14.0, 14.5, 15.0],
+        [2.2, 3.2, 4.0, 4.8, 5.3, 5.8],
+    ))
+    return max(1, int(round(notes_per_sec * duration_sec)))
+
+
+@torch.inference_mode()
+def generate_chart(
+    model: MaiGenerator,
+    audio_tokens: torch.Tensor,
+    bpm: float,
+    difficulty: int,
+    level_value: float,
+    genre: int = 0,
+    max_tokens: int = 4096,
+    temperature: float = 0.8,
+    top_k: int = 50,
+    top_p: float = 0.95,
+    eos_suppress: float = 1.5,
+    min_division: int = 4,
+    max_duration_beats: float = 4.0,
+    max_sim: int = 2,
+    sim_penalty: float = 1.0,
+    repetition_penalty: float = 1.0,
+    rest_streak_penalty: float = 0.0,
+    rest_streak_threshold: int = 4,
+    target_notes: int | None = None,
+    density_bias: float = 0.0,
+    break_penalty: float = 0.0,
+    break_window: int = 64,
+    max_break_ratio: float = 0.08,
+    type_biases: dict[str, float] | None = None,
+    config_density_bias: float = 0.0,
+    config_density_include_break: bool = False,
+) -> list[int]:
+    """
+    Autoregressive chart generation from audio (with KV-cache for speed).
+
+    Args:
+        model: Trained MaiGenerator.
+        audio_tokens: [1, T_aud] — EnCodec tokens.
+        bpm: BPM value.
+        difficulty: 0=BASIC..4=ReMASTER.
+        level_value: e.g. 12.4.
+        genre: Genre index.
+        max_tokens: Max tokens to generate.
+        temperature: Sampling temperature.
+        top_k: Top-K filter.
+        top_p: Nucleus filter.
+
+    Returns:
+        List of chart token IDs (including BOS, metadata header, and EOS).
+    """
+    model.eval()
+    device = next(model.parameters()).device
+    param_dtype = next(model.parameters()).dtype
+
+    if audio_tokens.dim() == 1:
+        audio_tokens = audio_tokens.unsqueeze(0)
+    audio_tokens = audio_tokens.to(device)
+
+    # ── Encode audio (once) ──
+    max_len = model.audio_pos_embed.num_embeddings  # 32768
+    if audio_tokens.shape[1] > max_len:
+        audio_tokens = audio_tokens[:, :max_len]
+    T_aud = audio_tokens.shape[1]
+    aud = model.audio_embed(audio_tokens)
+    aud = aud + model.audio_pos_embed(
+        torch.arange(T_aud, device=device).unsqueeze(0))
+
+    # Downsample audio (same as training forward)
+    if model.audio_downsample > 1:
+        aud = aud.transpose(1, 2)
+        aud = model.audio_down(aud)
+        aud = aud.transpose(1, 2)
+        T_aud = aud.shape[1]
+
+    aud = model.dropout(aud)
+    aud_pos = torch.arange(T_aud, device=device, dtype=torch.float32)
+    for blk in model.audio_encoder:
+        aud = blk(aud, positions=aud_pos)
+    aud = model.enc_norm(aud)
+
+    # ── Conditions ──
+    bpm_t = torch.tensor([[bpm]], device=device, dtype=param_dtype)
+    diff_t = torch.tensor([[difficulty]], device=device)
+    level_t = torch.tensor([[level_value]], device=device, dtype=param_dtype)
+    genre_t = torch.tensor([[genre]], device=device)
+
+    cond = (model.bpm_proj(bpm_t).unsqueeze(1) +
+            model.diff_embed(diff_t.squeeze(-1)).unsqueeze(1) +
+            model.level_proj(level_t).unsqueeze(1) +
+            model.genre_embed(genre_t.squeeze(-1)).unsqueeze(1))
+
+    diff_vec = model.diff_embed(diff_t.squeeze(-1))  # [1, d_model]
+
+    # Total beats
+    audio_duration_sec = audio_tokens.shape[1] / 2 / AUDIO_FRAME_RATE
+    total_beats = audio_duration_sec * bpm / 60.0
+    start_offset_beats = 2.0 * max(bpm, 30.0) / 60.0
+    print(f"Track: {audio_duration_sec:.0f}s = {total_beats:.0f} beats @ {bpm} BPM  (start at beat {start_offset_beats:.1f})")
+    if density_bias > 0 and target_notes is None:
+        target_notes = _auto_target_notes(level_value, audio_duration_sec)
+    if density_bias > 0 and target_notes:
+        print(f"Density target: {target_notes} notes  bias={density_bias}")
+
+    # ── Init KV-cache ──
+    model.init_kv_cache(batch_size=1, max_len=max_tokens + 1, enc_out=aud)
+
+    # ── Prefill with BOS token ──
+    tokens = [BOS]
+    pos_cache = [0.0]
+
+    tok_t = torch.tensor([[BOS]], device=device)
+    emb = model.chart_embed(tok_t) + cond  # cond adds batch, tok_t is [1,1]
+    chart_pos = torch.tensor([[0.0]], device=device, dtype=torch.float32)
+
+    x = emb
+    for blk in model.chart_decoder:
+        x = blk(x, enc_out=aud, self_positions=chart_pos,
+                 diff_emb=diff_vec, use_cache=True)
+
+    # ── Autoregressive loop (KV-cached, one token at a time) ──
+    decode_state = DecodeState(
+        bpm=bpm,
+        total_beats=total_beats,
+        start_offset_beats=start_offset_beats,
+        min_division=min_division,
+        max_sim=max_sim,
+        max_duration_beats=max_duration_beats,
+    )
+    note_bonus_ids = _density_note_ids(model.chart_vocab_size, device, include_break=False) if density_bias > 0 else None
+    break_ids = _break_ids(model.chart_vocab_size, device) if break_penalty > 0 else None
+    type_bias = _type_bias_vector(model.chart_vocab_size, device, type_biases or {})
+    config_density_bias_vec = _config_density_bias_vector(
+        model.chart_vocab_size, device, config_density_bias,
+        include_break=config_density_include_break,
+    )
+    if type_bias is not None:
+        print("Type bias:", ", ".join(f"{k}={v:+.2f}" for k, v in (type_biases or {}).items() if abs(v) > 1e-9))
+    if config_density_bias_vec is not None:
+        print(f"Config density bias: {config_density_bias:+.2f} "
+              f"(include_break={config_density_include_break})")
+    generated_notes = 0
+    generated_breaks = 0
+
+    pbar = tqdm(total=max_tokens, desc="Generating", unit="tok")
+    for step_idx in range(max_tokens):
+        x_last = model.dec_norm(x[:, -1:, :])
+        logits = model.output_head(x_last).squeeze(0).squeeze(0)
+
+        # ── Constraints ──
+        logits[EOS] = float("-inf")
+
+        if decode_state.current_beat >= total_beats:
+            pbar.update(1); pbar.set_postfix(tokens=len(tokens)+1, beat=f"{decode_state.current_beat:.0f}/{total_beats:.0f}", end="track")
+            tokens.append(EOS)
+            break
+
+        logits = decode_state.apply_logits_mask(logits)
+        if type_bias is not None:
+            logits = logits + type_bias.to(dtype=logits.dtype)
+        if config_density_bias_vec is not None:
+            logits = logits + config_density_bias_vec.to(dtype=logits.dtype)
+        if decode_state.can_start_sim() and SIM_BEG < logits.numel():
+            logits[SIM_BEG] -= sim_penalty
+        if not torch.isfinite(logits).any():
+            pbar.update(1); pbar.set_postfix(tokens=len(tokens)+1, beat=f"{decode_state.current_beat:.0f}/{total_beats:.0f}", end="no_legal")
+            tokens.append(EOS)
+            break
+
+        # ── Rest-streak penalty: if last N notes are all rests, suppress RST ──
+        if rest_streak_penalty > 0 and len(tokens) >= rest_streak_threshold:
+            recent = tokens[-rest_streak_threshold:]
+            # Check if recent tokens are all RST (not counting structural tokens like DUR/SIM etc.)
+            if all(t == RST for t in recent):
+                logits[RST] -= rest_streak_penalty
+
+        if density_bias > 0 and target_notes and note_bonus_ids is not None and note_bonus_ids.numel() > 0:
+            progress = max(0.0, decode_state.current_beat - start_offset_beats)
+            denom = max(1e-6, total_beats - start_offset_beats)
+            expected = float(target_notes) * min(1.0, progress / denom)
+            deficit = max(0.0, expected - generated_notes) / max(expected, 1.0)
+            if deficit > 0:
+                bonus = density_bias * deficit
+                logits[note_bonus_ids] += bonus
+                if RST < logits.numel():
+                    logits[RST] -= bonus
+
+        if break_penalty > 0 and break_ids is not None and break_ids.numel() > 0:
+            recent = tokens[-break_window:] if break_window > 0 else tokens
+            recent_notes = sum(_token_note_count(t) for t in recent)
+            recent_breaks = sum(_token_note_count(t) for t in recent if _token_has_break(t))
+            global_ratio = generated_breaks / max(generated_notes, 1)
+            recent_ratio = recent_breaks / max(recent_notes, 1)
+            excess = max(global_ratio - max_break_ratio, recent_ratio - max_break_ratio, 0.0)
+            if excess > 0:
+                logits[break_ids] -= break_penalty * (1.0 + 8.0 * excess)
+
+        next_tok = sample_token(logits, temperature, top_k, top_p,
+                                repetition_penalty=repetition_penalty,
+                                prev_tokens=tokens)
+        tokens.append(next_tok)
+        note_count = _token_note_count(next_tok)
+        generated_notes += note_count
+        if _token_has_break(next_tok):
+            generated_breaks += note_count
+
+        # ── Snap duration ──
+        if len(tokens) >= 4 and tokens[-4] not in (DUR_TOKEN, PAD) and tokens[-3] == DUR_TOKEN:
+            raw_num = ID_TO_DUR_NUM.get(tokens[-2], tokens[-2])
+            raw_den = ID_TO_DUR_DEN.get(next_tok, max(next_tok, 1))
+            beat, subdiv = snap_duration(raw_num, raw_den,
+                                          decode_state.remaining_duration_beats(),
+                                          decode_state.max_duration_beats)
+            tokens[-2], tokens[-1] = DUR_NUM_TO_ID[beat], DUR_DEN_TO_ID[subdiv]
+            next_tok = tokens[-1]
+
+        # ── Record position ──
+        bpm_v = max(bpm, 30.0)
+        time_sec = decode_state.current_beat * 60.0 / bpm_v
+        new_pos = time_sec * AUDIO_FRAME_RATE / model.audio_downsample
+        pos_cache.append(new_pos)
+
+        decode_state.step(next_tok)
+
+        # ── Incremental decoder step (KV-cached) ──
+        tok_t = torch.tensor([[next_tok]], device=device)
+        emb = model.chart_embed(tok_t) + cond
+        chart_pos = torch.tensor([[new_pos]], device=device, dtype=torch.float32)
+
+        x = emb
+        for blk in model.chart_decoder:
+            x = blk(x, enc_out=aud, self_positions=chart_pos,
+                     diff_emb=diff_vec, use_cache=True)
+
+        pbar.update(1)
+        if len(tokens) % 200 == 0:
+            pct = min(100, decode_state.current_beat / total_beats * 100)
+            pbar.set_postfix(beat=f"{decode_state.current_beat:.0f}/{total_beats:.0f}",
+                             pct=f"{pct:.0f}%", div=f"{decode_state.div_value}")
+    pbar.close()
+
+    tokens = sanitize_sim_tokens(tokens, max_sim=max_sim)
+    return clamp_duration_tokens(tokens, total_beats, start_offset_beats, min_division)
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# CLI
+# ═══════════════════════════════════════════════════════════════════════
+
+DIFF_MAP = {"BASIC": 0, "ADVANCED": 1, "EXPERT": 2, "MASTER": 3, "REMASTER": 4}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate maimai chart from audio")
+    parser.add_argument("--checkpoint", type=str, required=True,
+                        help="Path to model checkpoint")
+    parser.add_argument("--audio", type=str, default=None,
+                        help="Path to audio file (mp3/wav)")
+    parser.add_argument("--song-id", type=str, default=None,
+                        help="Song ID in datasets/ (uses datasets/<id>/track.mp3)")
+    parser.add_argument("--bpm", type=float, default=None,
+                        help="BPM (auto-detect from maidata.txt if --song-id)")
+    parser.add_argument("--diff", type=str, default="MASTER",
+                        choices=list(DIFF_MAP.keys()))
+    parser.add_argument("--level", type=float, default=12.0,
+                        help="Target level value")
+    parser.add_argument("--genre", type=int, default=0)
+    parser.add_argument("--max-tokens", type=int, default=4096)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top-k", type=int, default=50)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--eos-suppress", type=float, default=1.5,
+                        help="Subtract from EOS logit (0=off)")
+    parser.add_argument("--min-division", type=int, default=4,
+                        help="Minimum generated beat division; blocks div_1/div_2 by default")
+    parser.add_argument("--max-duration-beats", type=float, default=4.0,
+                        help="Maximum generated hold/slide duration in beats")
+    parser.add_argument("--max-sim", type=int, default=2,
+                        help="Maximum simultaneous press count")
+    parser.add_argument("--sim-penalty", type=float, default=1.0,
+                        help="Subtract from SIM_BEG logit before constrained sampling")
+    parser.add_argument("--rep-penalty", type=float, default=1.0,
+                        help="Repetition penalty (>1.0=penalize repeats, 1.0=off)")
+    parser.add_argument("--rest-streak-penalty", type=float, default=0.0,
+                        help="Penalize RST when last N tokens are all rests (try 2.0-5.0)")
+    parser.add_argument("--rest-streak-threshold", type=int, default=4,
+                        help="Consecutive RSTs before applying rest-streak-penalty")
+    parser.add_argument("--target-notes", type=int, default=None,
+                        help="Optional soft target note count for density-guided sampling")
+    parser.add_argument("--density-bias", type=float, default=0.0,
+                        help="Softly boost note/config tokens when generated density lags target")
+    parser.add_argument("--break-penalty", type=float, default=1.5,
+                        help="Type-level penalty when BREAK/绝赞 ratio is too high")
+    parser.add_argument("--break-window", type=int, default=64,
+                        help="Recent token window used by BREAK ratio penalty")
+    parser.add_argument("--max-break-ratio", type=float, default=0.08,
+                        help="Soft maximum BREAK/绝赞 ratio before penalty kicks in")
+    parser.add_argument("--tap-bias", type=float, default=0.0,
+                        help="Logit bias for TAP config/tokens (+ increases, - decreases)")
+    parser.add_argument("--hold-bias", type=float, default=0.0,
+                        help="Logit bias for HOLD config/tokens")
+    parser.add_argument("--slide-bias", type=float, default=0.0,
+                        help="Logit bias for SLIDE config/tokens")
+    parser.add_argument("--break-bias", type=float, default=0.0,
+                        help="Logit bias for BREAK/绝赞 config/tokens")
+    parser.add_argument("--touch-bias", type=float, default=0.0,
+                        help="Logit bias for TOUCH config/tokens")
+    parser.add_argument("--rest-bias", type=float, default=0.0,
+                        help="Logit bias for REST tokens")
+    parser.add_argument("--sim-bias", type=float, default=0.0,
+                        help="Logit bias for simultaneous group start")
+    parser.add_argument("--config-density-bias", type=float, default=0.0,
+                        help="Bias dense config tokens by extra notes per time slot")
+    parser.add_argument("--config-density-include-break", action="store_true",
+                        help="Allow config-density-bias to boost BREAK/绝赞 configs too")
+    parser.add_argument("--output", type=str, default=None,
+                        help="Output path for generated maidata.txt")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--precision", type=str, default="float32",
+                        choices=["float32", "float16", "bfloat16"],
+                        help="Model precision (float16/bfloat16 ≈2x faster on supported GPUs)")
+    parser.add_argument("--compile", action="store_true",
+                        help="Use torch.compile for faster execution (PyTorch 2.0+)")
+    args = parser.parse_args()
+
+    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    # ── Load model ──
+    ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False)
+    config = ckpt.get("config", {})
+    if "config_vocab" in config:
+        load_config_vocab(config["config_vocab"])
+        print(f"Loaded config vocab: total={MaiChartTokenizer.vocab_size}")
+    model = MaiGenerator(
+        d_model=config.get("d_model", 512),
+        enc_layers=config.get("enc_layers", 6),
+        dec_layers=config.get("dec_layers", 8),
+        heads=config.get("heads", 8),
+        d_ff=config.get("d_ff", 2048),
+        use_moe=config.get("use_moe", True),
+        n_experts=config.get("n_experts", 6),
+        moe_layers=config.get("moe_layers", None),
+    ).to(device)
+    model_state = model.state_dict()
+    filtered_state = {
+        k: v for k, v in ckpt["model_state"].items()
+        if k in model_state and tuple(v.shape) == tuple(model_state[k].shape)
+    }
+    skipped = len(ckpt["model_state"]) - len(filtered_state)
+    incompatible = model.load_state_dict(filtered_state, strict=False)
+    if skipped:
+        print(f"Model warm-start: skipped {skipped} shape-mismatched params")
+    if incompatible.missing_keys:
+        print(f"Model warm-start: initialized {len(incompatible.missing_keys)} new params")
+    if incompatible.unexpected_keys:
+        print(f"Model warm-start: ignored {len(incompatible.unexpected_keys)} checkpoint params")
+    model.eval()
+    print(f"Model loaded (epoch {ckpt.get('epoch', '?')})")
+
+    # ── Precision ──
+    if args.precision != "float32":
+        dt = {"float16": torch.float16, "bfloat16": torch.bfloat16}[args.precision]
+        model = model.to(dtype=dt)
+        print(f"Precision: {args.precision}")
+
+    # ── torch.compile ──
+    if args.compile:
+        try:
+            model = torch.compile(model, mode="reduce-overhead")
+            print("Model compiled with torch.compile (reduce-overhead)")
+        except Exception as e:
+            print(f"torch.compile failed: {e}, continuing without compile")
+
+    # ── Load / tokenize audio ──
+    audio_tok = MaiTrackTokenizer(n_layers=2, device=str(device))
+
+    if args.song_id:
+        audio_path = Path("datasets") / args.song_id / "track.mp3"
+        if not audio_path.exists():
+            print(f"ERROR: {audio_path} not found")
+            return
+        print(f"Audio: {audio_path}")
+        audio_tokens = audio_tok.encode(str(audio_path), add_bos=False, add_eos=False)
+
+        # Auto-detect BPM from maidata
+        if args.bpm is None:
+            mdata = Path("datasets") / args.song_id / "maidata.txt"
+            if mdata.exists():
+                song = parse_maidata_file(str(mdata))
+                bpm = song.bpm if song.bpm > 0 else 150.0
+                print(f"Auto-detected BPM: {bpm}")
+            else:
+                bpm = 150.0
+        else:
+            bpm = args.bpm
+    elif args.audio:
+        print(f"Audio: {args.audio}")
+        audio_tokens = audio_tok.encode(args.audio, add_bos=False, add_eos=False)
+        bpm = args.bpm or 150.0
+    else:
+        print("ERROR: need --audio or --song-id")
+        return
+
+    print(f"Audio tokens: {len(audio_tokens)}")
+
+    # ── Generate ──
+    difficulty = DIFF_MAP[args.diff.upper()]
+    print(f"Params: BPM={bpm}, Diff={args.diff}({difficulty}), Level={args.level}")
+    print(f"Generating (max {args.max_tokens} tokens)...")
+
+    audio_t = torch.tensor([audio_tokens], dtype=torch.long)
+    chart_tokens = generate_chart(
+        model, audio_t, bpm=bpm, difficulty=difficulty,
+        level_value=args.level, genre=args.genre,
+        max_tokens=args.max_tokens,
+        temperature=args.temperature, top_k=args.top_k, top_p=args.top_p,
+        eos_suppress=args.eos_suppress,
+        min_division=args.min_division,
+        max_duration_beats=args.max_duration_beats,
+        max_sim=args.max_sim,
+        sim_penalty=args.sim_penalty,
+        repetition_penalty=args.rep_penalty,
+        rest_streak_penalty=args.rest_streak_penalty,
+        rest_streak_threshold=args.rest_streak_threshold,
+        target_notes=args.target_notes,
+        density_bias=args.density_bias,
+        break_penalty=args.break_penalty,
+        break_window=args.break_window,
+        max_break_ratio=args.max_break_ratio,
+        type_biases={
+            "tap": args.tap_bias,
+            "hold": args.hold_bias,
+            "slide": args.slide_bias,
+            "break": args.break_bias,
+            "touch": args.touch_bias,
+            "rest": args.rest_bias,
+            "sim": args.sim_bias,
+        },
+        config_density_bias=args.config_density_bias,
+        config_density_include_break=args.config_density_include_break,
+    )
+
+    print(f"Generated {len(chart_tokens)} tokens")
+    audio_duration_sec = len(audio_tokens) / 2 / AUDIO_FRAME_RATE
+    total_beats = audio_duration_sec * bpm / 60.0
+    start_offset_beats = 2.0 * max(bpm, 30.0) / 60.0
+    dur_stats = duration_report(chart_tokens, total_beats, start_offset_beats, args.min_division)
+    print(f"Duration stats: count={dur_stats['durations']} "
+          f"max={dur_stats['max_beats']:.3f} beats overrun={dur_stats['overrun']}")
+    sim_stats = simultaneous_report(chart_tokens, args.max_sim)
+    print(f"Sim stats: groups={sim_stats['groups']} "
+          f"max_count={sim_stats['max_count']} over_limit={sim_stats['over_limit']} "
+          f"first_bad={sim_stats['first_bad']}")
+
+    # ── Decode to human-readable ──
+    chart_tok = MaiChartTokenizer()
+    print(f"\nTokens: {chart_tok.tokens_to_str(chart_tokens, 60)}")
+
+    # Decode to chart
+    decoded = chart_tok.decode(chart_tokens)
+    print(f"\nChart stats: TAP={decoded.tap_count} HOLD={decoded.hold_count} "
+          f"SLIDE={decoded.slide_count} BREAK={decoded.break_count} "
+          f"TOUCH={decoded.touch_count} TOTAL={decoded.note_count}")
+    print(f"Break ratio: {decoded.break_count / max(decoded.note_count, 1):.2%}")
+    print(f"Density: {decoded.note_count / max(audio_duration_sec, 1e-6):.2f} notes/s, "
+          f"{decoded.note_count / max(total_beats, 1e-6):.2f} notes/beat")
+
+    # ── Save ──
+    if args.output:
+        out_path = Path(args.output)
+        # Convert tokens → proper maidata chart text
+        chart_text = tokens_to_maitext(chart_tokens, bpm)
+        lines = [f"&title=Generated Chart",
+                 f"&wholebpm={bpm}",
+                 f"&artist=MaiGenerator",
+                 f"&lv_{difficulty+1}={args.level}",
+                 f"&inote_{difficulty+1}="]
+        lines.append(chart_text)
+        out_path.write_text("\n".join(lines).rstrip("\n"), encoding="utf-8")
+        print(f"\nSaved to {out_path}")
+
+
+if __name__ == "__main__":
+    main()

mai_parser/__init__.py

@@ -0,0 +1,55 @@
+"""
+mai_parser - Maimai (舞萌) chart data parser for ML training.
+
+Usage:
+    from mai_parser import SongDataset, parse_maidata_file
+
+    # Load entire dataset
+    ds = SongDataset("datasets").load()
+    print(ds.stats())
+
+    # Parse single file
+    song = parse_maidata_file("datasets/10/maidata.txt")
+    print(song.title_clean, song.artist, song.bpm)
+
+    # Convert chart to tensor for training
+    chart = song.get_chart(Difficulty.MASTER)
+    tensor = chart_to_tensor(chart)
+"""
+
+from .models import Cabinet, Chart, Difficulty, Song, TouchNote
+from .parser import parse_maidata, parse_maidata_file
+from .dataset import (
+    ChartDataset,
+    SongDataset,
+    chart_to_sequence,
+    chart_to_tensor,
+    note_to_vector,
+    NOTE_TYPE,
+    NOTE_TYPE_INV,
+    NUM_POSITIONS,
+    HAS_TORCH,
+)
+
+__all__ = [
+    # Models
+    "Cabinet",
+    "Chart",
+    "Difficulty",
+    "Song",
+    "TouchNote",
+    # Parser
+    "parse_maidata",
+    "parse_maidata_file",
+    # Dataset
+    "ChartDataset",
+    "SongDataset",
+    "chart_to_sequence",
+    "chart_to_tensor",
+    "note_to_vector",
+    # Constants
+    "NOTE_TYPE",
+    "NOTE_TYPE_INV",
+    "NUM_POSITIONS",
+    "HAS_TORCH",
+]

mai_parser/dataset.py

@@ -0,0 +1,455 @@
+"""
+Dataset loader for maimai training data.
+
+Provides:
+- SongDataset: scans datasets/ folder, parses all songs
+- ChartDataset: PyTorch Dataset for chart note sequences
+- Utility functions for note-to-tensor conversion
+"""
+
+from __future__ import annotations
+
+import json
+import pickle
+from pathlib import Path
+from typing import Iterator, Optional, Union
+
+from .models import Cabinet, Chart, Difficulty, Song, TouchNote
+from .parser import parse_maidata_file
+
+# Try importing PyTorch (optional)
+try:
+    import torch
+    from torch.utils.data import Dataset as TorchDataset
+    HAS_TORCH = True
+except (ImportError, OSError):
+    HAS_TORCH = False
+
+
+class SongDataset:
+    """
+    A collection of parsed Song objects from the datasets/ directory.
+
+    Usage:
+        ds = SongDataset("datasets")
+        ds.scan()              # discover all songs
+        ds.load(use_cache=True)  # parse all songs (with caching)
+        print(len(ds))          # total songs
+        song = ds[0]            # first song
+    """
+
+    def __init__(self, root_dir: str | Path):
+        self.root_dir = Path(root_dir)
+        self.songs: list[Song] = []
+        self._by_id: dict[str, Song] = {}
+        self._cache_path = self.root_dir / ".mai_parser_cache.pkl"
+
+    def scan(self) -> list[Path]:
+        """
+        Scan the dataset directory for all song folders (containing maidata.txt).
+
+        Returns list of found song folder paths.
+        """
+        folders = []
+        for item in sorted(self.root_dir.iterdir()):
+            if item.is_dir() and (item / "maidata.txt").exists():
+                folders.append(item)
+        return folders
+
+    def load(self, use_cache: bool = True, max_workers: int = 8,
+             progress: bool = True) -> "SongDataset":
+        """
+        Parse all songs from discovered folders.
+
+        Args:
+            use_cache: If True, try loading from cache first; save after parsing.
+            max_workers: Number of threads for parallel parsing.
+            progress: Print progress information.
+        """
+        if use_cache and self._cache_path.exists():
+            try:
+                with open(self._cache_path, "rb") as f:
+                    self.songs = pickle.load(f)
+                self._build_index()
+                if progress:
+                    print(f"Loaded {len(self.songs)} songs from cache.")
+                return self
+            except (pickle.UnpicklingError, EOFError, KeyError) as e:
+                if progress:
+                    print(f"Cache invalid ({e}), re-parsing...")
+
+        folders = self.scan()
+        if progress:
+            print(f"Found {len(folders)} song folders. Parsing...")
+
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+
+        songs: list[Song] = []
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(parse_maidata_file, f / "maidata.txt"): f
+                for f in folders
+            }
+            for i, future in enumerate(as_completed(futures)):
+                folder = futures[future]
+                try:
+                    song = future.result()
+                    song.maidata_path = str(folder / "maidata.txt")
+                    song.audio_path = str(folder / "track.mp3")
+                    songs.append(song)
+                except Exception as e:
+                    if progress:
+                        print(f"  [WARN] Failed to parse {folder.name}: {e}")
+
+                if progress and (i + 1) % 200 == 0:
+                    print(f"  ... parsed {i + 1}/{len(folders)}")
+
+        # Sort by song_id numerically if possible
+        songs.sort(key=lambda s: _sort_key(s.song_id))
+
+        self.songs = songs
+        self._build_index()
+
+        if use_cache:
+            try:
+                with open(self._cache_path, "wb") as f:
+                    pickle.dump(self.songs, f, protocol=pickle.HIGHEST_PROTOCOL)
+                if progress:
+                    print(f"Saved cache to {self._cache_path}")
+            except OSError as e:
+                if progress:
+                    print(f"  [WARN] Could not save cache: {e}")
+
+        if progress:
+            print(f"Loaded {len(self.songs)} songs.")
+        return self
+
+    def _build_index(self) -> None:
+        self._by_id = {s.song_id: s for s in self.songs}
+
+    def __len__(self) -> int:
+        return len(self.songs)
+
+    def __getitem__(self, idx: int) -> Song:
+        return self.songs[idx]
+
+    def __iter__(self) -> Iterator[Song]:
+        return iter(self.songs)
+
+    def get_by_id(self, song_id: str) -> Song | None:
+        """Look up a song by its folder name / ID."""
+        return self._by_id.get(song_id)
+
+    def filter(
+        self,
+        cabinet: Cabinet | None = None,
+        min_level: float | None = None,
+        max_level: float | None = None,
+        difficulty: Difficulty | None = None,
+        genre: str | None = None,
+        has_audio: bool = False,
+    ) -> list[Song]:
+        """
+        Filter songs by various criteria.
+
+        Args:
+            cabinet: Filter by cabinet type (SD/DX).
+            min_level: Minimum chart level.
+            max_level: Maximum chart level.
+            difficulty: Only songs with this difficulty.
+            genre: Filter by genre name (partial match).
+            has_audio: Only songs where track.mp3 exists.
+        """
+        result = []
+        for song in self.songs:
+            if cabinet is not None and song.cabinet != cabinet:
+                continue
+            if genre is not None and genre.lower() not in song.genre.lower():
+                continue
+            if has_audio:
+                audio_path = self.root_dir / song.song_id / "track.mp3"
+                if not audio_path.exists():
+                    continue
+            if difficulty is not None:
+                chart = song.get_chart(difficulty)
+                if chart is None:
+                    continue
+                if min_level is not None and chart.level_value < min_level:
+                    continue
+                if max_level is not None and chart.level_value > max_level:
+                    continue
+            result.append(song)
+        return result
+
+    def stats(self) -> dict:
+        """Return summary statistics of the dataset."""
+        total = len(self.songs)
+        sd_count = sum(1 for s in self.songs if s.cabinet == Cabinet.SD)
+        dx_count = sum(1 for s in self.songs if s.cabinet == Cabinet.DX)
+        unknown_cab = total - sd_count - dx_count
+        total_charts = sum(len(s.charts) for s in self.songs)
+        genres = list(set(s.genre for s in self.songs if s.genre))
+
+        # Level distribution per difficulty
+        level_dist: dict[str, list[float]] = {}
+        for song in self.songs:
+            for idx, chart in song.charts.items():
+                key = f"lv_{idx}"
+                if key not in level_dist:
+                    level_dist[key] = []
+                level_dist[key].append(chart.level_value)
+
+        return {
+            "total_songs": total,
+            "sd_songs": sd_count,
+            "dx_songs": dx_count,
+            "unknown_cabinet": unknown_cab,
+            "total_charts": total_charts,
+            "genres": sorted(genres),
+            "level_distribution": {
+                k: {
+                    "min": min(v) if v else 0,
+                    "max": max(v) if v else 0,
+                    "avg": sum(v) / len(v) if v else 0,
+                }
+                for k, v in level_dist.items()
+            },
+        }
+
+
+def _sort_key(song_id: str) -> tuple:
+    """Sort: numeric IDs first, then strings."""
+    try:
+        return (0, int(song_id))
+    except ValueError:
+        return (1, song_id)
+
+
+# ─── Note to Tensor Conversion ──────────────────────────────────────────────
+
+# Note type encoding for training
+NOTE_TYPE = {
+    "rest": 0,
+    "tap": 1,
+    "hold": 2,
+    "slide": 3,
+    "break": 4,
+    "touch": 5,
+    "end": 6,
+}
+
+# Reverse mapping
+NOTE_TYPE_INV = {v: k for k, v in NOTE_TYPE.items()}
+
+# Number of button positions (1-8)
+NUM_POSITIONS = 8
+
+
+def note_to_vector(note: TouchNote, num_positions: int = NUM_POSITIONS) -> list[float]:
+    """
+    Convert a single TouchNote to a feature vector.
+
+    Format (fixed-length):
+      [type_rest, type_tap, type_hold, type_slide, type_break, type_touch, type_end,
+       pos_1, pos_2, ..., pos_N,
+       is_simultaneous, has_slide_path_len_2, has_slide_path_len_3, has_slide_path_len_4,
+       hold_beat, hold_subdiv,
+       beat_div_normalized]
+
+    Returns a flat list of floats.
+    """
+    vec = [0.0] * (7 + num_positions + 4 + 2 + 1)
+
+    # Note type (one-hot)
+    if note.is_end:
+        vec[6] = 1.0
+    elif note.is_rest:
+        vec[0] = 1.0
+    elif note.is_touch:
+        vec[5] = 1.0
+    elif note.is_break:
+        vec[4] = 1.0
+    elif note.is_hold:
+        vec[2] = 1.0
+    elif note.is_slide:
+        vec[3] = 1.0
+    else:
+        vec[1] = 1.0  # tap
+
+    # Button positions (multi-hot)
+    offset = 7
+    for pos in note.positions:
+        if 1 <= pos <= num_positions:
+            vec[offset + pos - 1] = 1.0
+
+    # Slide path length flag
+    offset += num_positions
+    sl = len(note.slide_path)
+    if sl >= 2:
+        vec[offset] = 1.0
+    if sl >= 3:
+        vec[offset + 1] = 1.0
+    if sl >= 4:
+        vec[offset + 2] = 1.0
+    vec[offset + 3] = 1.0 if note.is_simultaneous else 0.0
+
+    # Hold duration
+    offset += 4
+    if note.hold_duration:
+        vec[offset] = float(note.hold_duration[0])
+        vec[offset + 1] = float(note.hold_duration[1])
+
+    # Beat division (normalized)
+    vec[-1] = note.beat_div
+
+    return vec
+
+
+def chart_to_tensor(chart: Chart) -> "torch.Tensor":
+    """
+    Convert an entire chart's notes to a PyTorch tensor.
+
+    Returns tensor of shape (num_notes, feature_dim).
+    """
+    if not HAS_TORCH:
+        raise ImportError("PyTorch is required for tensor conversion. "
+                          "Install with: pip install torch")
+    vectors = [note_to_vector(n) for n in chart.notes]
+    return torch.tensor(vectors, dtype=torch.float32)
+
+
+def chart_to_sequence(chart: Chart) -> list[list[float]]:
+    """Convert a chart to a list of feature vectors (no PyTorch dependency)."""
+    return [note_to_vector(n) for n in chart.notes]
+
+
+# ─── PyTorch Dataset (if torch available) ───────────────────────────────────
+
+if HAS_TORCH:
+
+    class ChartDataset(TorchDataset):
+        """
+        PyTorch Dataset for maimai chart training.
+
+        Yields (note_sequence, metadata) pairs where:
+        - note_sequence is a tensor of shape (seq_len, feature_dim)
+        - metadata is a dict with song info
+
+        Usage:
+            from torch.utils.data import DataLoader
+            ds = ChartDataset(song_dataset, difficulty=Difficulty.MASTER)
+            loader = DataLoader(ds, batch_size=32, collate_fn=ds.collate_fn)
+        """
+
+        def __init__(
+            self,
+            song_dataset: SongDataset,
+            difficulty: Difficulty | None = None,
+            min_level: float | None = None,
+            max_level: float | None = None,
+            pad_to: int | None = None,
+            include_metadata: bool = True,
+        ):
+            self.song_dataset = song_dataset
+            self.difficulty = difficulty
+            self.min_level = min_level
+            self.max_level = max_level
+            self.pad_to = pad_to
+            self.include_metadata = include_metadata
+
+            # Build filtered list of (song_index, chart_index)
+            self._items: list[tuple[int, int]] = []
+            for si, song in enumerate(song_dataset.songs):
+                for idx, chart in song.charts.items():
+                    if difficulty is not None and idx != difficulty.value:
+                        continue
+                    if min_level is not None and chart.level_value < min_level:
+                        continue
+                    if max_level is not None and chart.level_value > max_level:
+                        continue
+                    self._items.append((si, idx))
+
+        def __len__(self) -> int:
+            return len(self._items)
+
+        def __getitem__(self, idx: int):
+            si, ci = self._items[idx]
+            song = self.song_dataset[si]
+            chart = song.charts[ci]
+
+            tensor = chart_to_tensor(chart)
+
+            if self.pad_to is not None:
+                if tensor.shape[0] < self.pad_to:
+                    pad = torch.zeros(
+                        self.pad_to - tensor.shape[0], tensor.shape[1]
+                    )
+                    tensor = torch.cat([tensor, pad], dim=0)
+
+            if self.include_metadata:
+                return tensor, {
+                    "song_id": song.song_id,
+                    "title": song.title_clean,
+                    "artist": song.artist,
+                    "bpm": song.bpm,
+                    "level": chart.level_value,
+                    "difficulty": chart.difficulty_index,
+                    "cabinet": song.cabinet.value,
+                    "note_count": chart.note_count,
+                }
+            return tensor
+
+        @staticmethod
+        def collate_fn(batch):
+            """Custom collate for variable-length sequences."""
+            if isinstance(batch[0], tuple):
+                sequences = [item[0] for item in batch]
+                metadata = [item[1] for item in batch]
+            else:
+                sequences = batch
+                metadata = None
+
+            # Pad sequences to max length in batch
+            lengths = torch.tensor([s.shape[0] for s in sequences])
+            max_len = lengths.max().item()
+            feat_dim = sequences[0].shape[1]
+
+            padded = torch.zeros(len(sequences), max_len, feat_dim)
+            for i, seq in enumerate(sequences):
+                padded[i, :seq.shape[0]] = seq
+
+            if metadata is not None:
+                return padded, lengths, metadata
+            return padded, lengths
+
+        def stats(self) -> dict:
+            """Return dataset statistics."""
+            levels = []
+            note_counts = []
+            for _, ci in self._items:
+                song_idx, chart_idx = self._items[0]
+            for si, ci in self._items:
+                chart = self.song_dataset[si].charts[ci]
+                levels.append(chart.level_value)
+                note_counts.append(chart.note_count)
+
+            return {
+                "total_charts": len(self._items),
+                "level_min": min(levels) if levels else 0,
+                "level_max": max(levels) if levels else 0,
+                "level_avg": sum(levels) / len(levels) if levels else 0,
+                "note_count_min": min(note_counts) if note_counts else 0,
+                "note_count_max": max(note_counts) if note_counts else 0,
+                "note_count_avg": sum(note_counts) / len(note_counts)
+                if note_counts else 0,
+            }
+
+else:
+
+    class ChartDataset:
+        """Stub when PyTorch is not installed."""
+        def __init__(self, *args, **kwargs):
+            raise ImportError(
+                "PyTorch is required for ChartDataset. "
+                "Install with: pip install torch"
+            )

mai_parser/models.py

@@ -0,0 +1,175 @@
+"""
+Data models for maimai chart parsing.
+
+Represents the full structure of a maimai song entry:
+- Song metadata (title, artist, BPM, version, etc.)
+- Difficulty configs (level, charter)
+- Chart note sequences
+"""
+
+from __future__ import annotations
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+
+class Cabinet(Enum):
+    """Cabinet / hardware type."""
+    SD = "SD"
+    DX = "DX"
+    UNKNOWN = "UNKNOWN"
+
+
+class Difficulty(Enum):
+    """Difficulty tiers (1-indexed matching &inote_N)."""
+    BASIC = 1       # lv_1 / inote_1
+    ADVANCED = 2    # lv_2 / inote_2
+    EXPERT = 3      # lv_3 / inote_3
+    MASTER = 4      # lv_4 / inote_4
+    ReMASTER = 5    # lv_5 / inote_5 → SD Re:MASTER is inote_6
+    # DX uses lv_6 for Re:MASTER (inote_6)
+    # Some charts use lv_7 / inote_7 for Ura/裏 or extra
+
+    @classmethod
+    def from_index(cls, idx: int) -> "Difficulty":
+        for d in cls:
+            if d.value == idx:
+                return d
+        return cls.ReMASTER  # fallback for 5+
+
+
+@dataclass
+class TouchNote:
+    """
+    A single parsed note/event.
+
+    Raw format examples:
+      {4}1       → beat_div=4, pos=1
+      {4}1/8     → beat_div=4, pos=1/8 (simultaneous press)
+      {4}1h[4:1] → beat_div=4, pos=1, modifier=hold, duration=4:1
+      {4}1b      → beat_div=4, pos=1, is_break=True
+      {4}3>6[4:1]→ beat_div=4, pos=3>6 (slide), modifier=slide, duration=4:1
+      {1}        → beat_div=1, rest
+      E          → end marker
+    """
+    beat_div: int = 4          # {4}, {8}, {16}, {1}, etc.
+    raw: str = ""              # original text
+    is_rest: bool = False
+    is_end: bool = False
+    is_break: bool = False     # 'b' suffix on note
+    is_star: bool = False      # '*' suffix (EX note)
+    is_slide: bool = False     # '>' or 'V' or other slide patterns
+    is_hold: bool = False      # 'h[...]' modifier
+    is_touch: bool = False     # touch panel notes (C, Cf, Ch, etc.)
+    is_simultaneous: bool = False  # contains '/' for multi-tap
+    positions: list[int] = field(default_factory=list)  # 1-8 button positions
+    touch_regions: list[str] = field(default_factory=list)  # e.g. ["C", "B7", "E8"]
+    hold_duration: tuple[int, int] | None = None  # (beat, subdiv) from h[beat:subdiv]
+    slide_path: list[int] = field(default_factory=list)  # slide positions
+    firework: bool = False     # 'x' suffix
+    ex_notes: list[str] = field(default_factory=list)  # EX modifiers
+
+
+@dataclass
+class Chart:
+    """
+    A single difficulty chart (one &inote_N block).
+
+    Contains the parsed note sequence and metadata about the chart.
+    """
+    difficulty_index: int         # 1-7, maps to &inote_N
+    difficulty: Difficulty        # enum value
+    level: str = ""               # e.g. "12.4", "13+", "14.6?"
+    level_value: float = 0.0      # parsed numeric level
+    is_plus: bool = False         # has '+' suffix
+    is_ura: bool = False          # has '?' suffix (裏譜面)
+    charter: str = ""             # &des_N, charter/description
+    notes: list[TouchNote] = field(default_factory=list)
+    # Computed stats
+    note_count: int = 0           # total non-rest notes
+    tap_count: int = 0
+    hold_count: int = 0
+    slide_count: int = 0
+    break_count: int = 0
+    touch_count: int = 0
+    total_beats: float = 0.0      # total chart duration in beats
+
+    def compute_stats(self) -> None:
+        """Recalculate computed statistics from notes."""
+        self.note_count = 0
+        self.tap_count = 0
+        self.hold_count = 0
+        self.slide_count = 0
+        self.break_count = 0
+        self.touch_count = 0
+
+        for n in self.notes:
+            if n.is_rest or n.is_end:
+                continue
+            self.note_count += 1
+            if n.is_touch:
+                self.touch_count += 1
+            elif n.is_break:
+                self.break_count += 1
+            elif n.is_hold:
+                self.hold_count += 1
+            elif n.is_slide:
+                self.slide_count += 1
+            else:
+                self.tap_count += 1
+
+
+@dataclass
+class Song:
+    """
+    Full parsed song entry — maps to one folder in datasets/.
+    """
+    # Identifiers
+    song_id: str = ""             # folder name (e.g. "10", "10021")
+    title: str = ""               # &title (raw, may include [SD] etc.)
+    title_clean: str = ""         # cleaned title without tags
+    artist: str = ""              # &artist
+    artist_id: int = 0            # &artistid
+    genre: str = ""               # &genre
+    genre_id: int = 0             # &genreid
+
+    # Music info
+    bpm: float = 0.0              # &wholebpm
+    first: float = 0.0            # &first (offset in seconds)
+
+    # Cabinet & version
+    cabinet: Cabinet = Cabinet.UNKNOWN
+    version: str = ""             # &version
+    short_id: int = 0             # &shortid
+
+    # Chart converter metadata
+    converter: str = ""            # &ChartConverter
+    converter_tool: str = ""       # &ChartConvertTool
+    converter_version: str = ""    # &ChartConvertToolVersion
+
+    # Raw header
+    description: str = ""          # &des
+
+    # Difficulty info (keyed by difficulty index 1-7)
+    levels: dict[int, str] = field(default_factory=dict)    # &lv_N
+    charters: dict[int, str] = field(default_factory=dict)  # &des_N
+
+    # Parsed charts (keyed by difficulty index 1-7)
+    charts: dict[int, Chart] = field(default_factory=dict)
+
+    # Flags
+    is_fulltouch: bool = False
+    is_full: bool = False          # [FULL] in title
+    is_utage: bool = False         # 宴会場 genre
+    tags: list[str] = field(default_factory=list)  # [SD], [DX], [宴], etc.
+
+    # File paths (relative to dataset root)
+    maidata_path: str = ""
+    audio_path: str = ""
+
+    def has_chart(self, diff: Difficulty) -> bool:
+        """Check if a difficulty level exists."""
+        return diff.value in self.charts
+
+    def get_chart(self, diff: Difficulty) -> Chart | None:
+        return self.charts.get(diff.value)

mai_parser/parser.py

@@ -0,0 +1,514 @@
+"""
+Parser for maidata.txt files (MaichartConverter format).
+
+Parses:
+1. Header key-value pairs (&key=value)
+2. Note sequence data (&inote_N= blocks)
+
+The note format is measure-based:
+  {beat_div}pos1,pos2,pos3,...
+
+Where each position can be:
+  - A number 1-8 (button position)
+  - Empty (rest)
+  - Number/Number (simultaneous tap)
+  - Number + modifiers (h=hold, b=break, x=firework, >=/</V/slide patterns)
+  - Touch region (C, B1-B8, E1-E8, A1-A8, D1-D8)
+  - End marker: "E"
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Optional
+
+from .models import Cabinet, Chart, Difficulty, Song, TouchNote
+
+# ─── Regex patterns for note parsing ────────────────────────────────────────
+
+# Matches {beat_div} at the start of a measure
+_RE_BEAT_DIV = re.compile(r"\{(\d+)\}")
+
+# Matches inline BPM changes: (200)
+_RE_BPM_CHANGE = re.compile(r"\((\d+(?:\.\d+)?)\)")
+
+# Matches hold duration: h[beat:subdiv]
+_RE_HOLD = re.compile(r"h\[(\d+):(\d+)\]")
+
+# Matches slide pattern with duration: >, <, -, V, v, ^
+# Examples: 3>6[4:1], 6<3[8:11], 3-6[4:1], 3v6[16:3], 5>2[8:11], 3>8-6v3[8:9]
+_RE_SLIDE_DUR = re.compile(r"([><Vv\-\^])(\d+)\[(\d+):(\d+)\]")
+
+# Matches simple slide (no duration): 3>6, 4-1
+_RE_SIMPLE_SLIDE = re.compile(r"(\d+)[><Vv\-\^](\d+)")
+
+# Matches break suffix: b or $b at end
+_RE_BREAK = re.compile(r"b$|\$b$")
+
+# Matches firework/EX: x suffix
+_RE_FIREWORK = re.compile(r"x$")
+
+# Matches touch note identifiers: C, Cf, Ch, B1-B8, E1-E8, A1-A8, D1-D8
+_RE_TOUCH = re.compile(r"^(C[fh]?|[ABDE][1-8])$")
+
+# Matches simultaneous tap: 1/8, 2/5/7
+_RE_SIMUL = re.compile(r"^(\d+)/(\d+)(?:/(\d+))?(?:/(\d+))?$")
+
+# Matches a number position (1-8)
+_RE_POS = re.compile(r"^\d+$")
+
+
+def parse_header_line(line: str) -> tuple[str, str]:
+    """Parse a single &key=value header line. Returns (key, value)."""
+    line = line.strip()
+    if line.startswith("&"):
+        line = line[1:]  # strip &
+    if "=" in line:
+        key, value = line.split("=", 1)
+        return key.strip(), value.strip()
+    return line.strip(), ""
+
+
+def _extract_bpm_changes(text: str) -> list[tuple[int, float]]:
+    """Extract inline BPM changes from a measure line.
+    Returns list of (position_in_measure, bpm_value).
+    """
+    changes = []
+    for m in _RE_BPM_CHANGE.finditer(text):
+        changes.append((m.start(), float(m.group(1))))
+    return changes
+
+
+def _parse_single_note(token: str, current_beat_div: int) -> TouchNote:
+    """
+    Parse a single note token (one comma-separated segment).
+
+    Examples:
+      ""        → rest
+      "E"       → end marker
+      "1"       → tap at position 1
+      "1/8"     → simultaneous tap at 1 and 8
+      "1b"      → break at position 1
+      "1x"      → firework at position 1
+      "1h[4:1]" → hold at position 1, duration 4:1
+      "3>6[4:1]"→ slide from 3 to 6, duration 4:1
+      "3-6"     → simple slide 3→6
+      "C"       → touch center
+      "B7/B6"   → touch simultaneous
+      "Ch[2:1]" → touch hold
+    """
+    note = TouchNote(beat_div=current_beat_div, raw=token)
+
+    # Empty = rest
+    if not token:
+        note.is_rest = True
+        return note
+
+    # End marker
+    if token.strip() == "E":
+        note.is_end = True
+        return note
+
+    # Work on the cleaned token (remove beat division, BPM)
+    t = token.strip()
+
+    # Remove inline BPM changes for parsing: (200), (150.5)
+    t = _RE_BPM_CHANGE.sub("", t).strip()
+    if not t:
+        note.is_rest = True
+        return note
+
+    # --- Check for touch notes ---
+    # Touch notes can have modifiers like Ch[2:1], C, B7/B6
+    # They use letters A-E + number, or just C
+    touch_parts = re.split(r"[/]", t)
+    all_touch = all(
+        _RE_TOUCH.match(re.sub(r"\[.*?\]", "", p).strip("hbfxqpb$Vv*^-><"))
+        for p in touch_parts
+    )
+
+    if all_touch and any(
+        re.match(r"^[A-E]", p.strip("hbfxqpb$Vv*^-><"))
+        for p in touch_parts
+    ):
+        note.is_touch = True
+        # Extract touch regions (strip modifiers)
+        for p in touch_parts:
+            clean = re.sub(r"\[.*?\]", "", p).strip("hbfxqpb$Vv*^-><")
+            if clean:
+                note.touch_regions.append(clean)
+        # Parse hold/touch-hold
+        hold_m = _RE_HOLD.search(t)
+        if hold_m:
+            note.is_hold = True
+            note.hold_duration = (int(hold_m.group(1)), int(hold_m.group(2)))
+        return note
+
+    # --- Parse button note ---
+
+    # Break suffix
+    if _RE_BREAK.search(t):
+        note.is_break = True
+        t = _RE_BREAK.sub("", t).rstrip("$")
+
+    # Firework suffix
+    if t.endswith("x"):
+        note.is_firework = True
+        t = t[:-1]
+
+    # Hold
+    hold_m = _RE_HOLD.search(t)
+    if hold_m:
+        note.is_hold = True
+        note.hold_duration = (int(hold_m.group(1)), int(hold_m.group(2)))
+        t = _RE_HOLD.sub("", t)
+
+    # Star/EX note (* suffix)
+    if t.endswith("*"):
+        note.is_star = True
+        t = t[:-1]
+
+    # Slide patterns
+    # Pattern with duration: 3>6[4:1], 3-6[8:11], 3v6[16:3], 3>8-6v3[8:9]
+    slide_m = _RE_SLIDE_DUR.search(t)
+    if slide_m:
+        note.is_slide = True
+        note.slide_path = _extract_slide_path(t)
+        hold_m2 = _RE_HOLD.search(t)
+        if not hold_m2:
+            # Duration is on slide: [beat:subdiv]
+            note.hold_duration = (int(slide_m.group(3)), int(slide_m.group(4)))
+        # Extract positions from slide
+        pos_match = re.findall(r"(\d+)", t)
+        note.positions = [int(p) for p in pos_match[:4]]
+        return note
+
+    # Simple slide: 3>6, 4-1, 8<5
+    simple_slide = _RE_SIMPLE_SLIDE.search(t)
+    if simple_slide:
+        note.is_slide = True
+        pos_match = re.findall(r"(\d+)", t)
+        note.positions = [int(p) for p in pos_match[:4]]
+        return note
+
+    # Single V (slide continuation / endpoint)
+    if re.match(r"^V\d*$", t):
+        note.is_slide = True
+        return note
+
+    # Simultaneous tap: 1/8, 2/6/7
+    simul_m = _RE_SIMUL.match(t)
+    if simul_m:
+        note.is_simultaneous = True
+        note.positions = [int(g) for g in simul_m.groups() if g is not None]
+        return note
+
+    # Simple tap: 1-8
+    if _RE_POS.match(t):
+        note.positions = [int(t)]
+        return note
+
+    # Special modifiers: 'q', 'p', 'w', 'z', '$' etc. — treat as tap if has number
+    num_match = re.findall(r"(\d+)", t)
+    if num_match:
+        note.positions = [int(p) for p in num_match[:4]]
+        return note
+
+    # Unknown / rest
+    note.is_rest = True
+    return note
+
+
+def _extract_slide_path(token: str) -> list[int]:
+    """Extract all position numbers from a slide pattern like '3>8-6v3[8:9]' → [3,8,6,3]."""
+    nums = re.findall(r"(?<![a-zA-Z\[:])(\d+)(?![a-zA-Z\]])", token)
+    return [int(n) for n in nums]
+
+
+def _parse_note_sequence(raw_text: str) -> list[TouchNote]:
+    """
+    Parse the full note sequence from an &inote_N block.
+
+    Handles:
+    - Beat division changes: {4}, {8}, {16}, etc.
+    - Inline BPM changes: (200)
+    - Measure-by-measure comma-separated notes
+    - End marker: E
+    """
+    notes: list[TouchNote] = []
+    current_div = 4  # default beat division
+
+    # Normalize: join continuation lines, split by newline
+    text = raw_text.strip()
+
+    for line in text.split("\n"):
+        line = line.strip()
+        if not line:
+            continue
+
+        # Check for beat division change at start of line
+        bd_match = _RE_BEAT_DIV.match(line)
+        if bd_match:
+            current_div = int(bd_match.group(1))
+            line = line[bd_match.end():]  # remove {beat_div} prefix
+
+        # If the line is empty after removing beat div, it's just a beat division
+        # announcement; subsequent lines use this division
+        if not line:
+            continue
+
+        # The line may start with an inline BPM change like (173){1},
+        bpm_match = _RE_BPM_CHANGE.match(line)
+        if bpm_match:
+            line = line[bpm_match.end():]
+            # Check for beat div inside the BPM line
+            bd2 = _RE_BEAT_DIV.match(line) if line else None
+            if bd2:
+                current_div = int(bd2.group(1))
+                line = line[bd2.end():]
+
+        if not line:
+            continue
+
+        # Split by commas for individual notes
+        tokens = line.split(",")
+
+        for token in tokens:
+            # Check for inline beat division change
+            bd_inline = _RE_BEAT_DIV.search(token) if token else None
+            local_div = current_div
+            if bd_inline:
+                local_div = int(bd_inline.group(1))
+                token = _RE_BEAT_DIV.sub("", token).strip()
+
+            note = _parse_single_note(token, local_div)
+            notes.append(note)
+
+            if note.is_end:
+                return notes  # stop at E marker
+
+    return notes
+
+
+def parse_level_value(level_str: str) -> tuple[float, bool, bool]:
+    """
+    Parse a level string.
+    Returns (numeric_value, is_plus, is_ura).
+
+    Examples:
+      "12.4"   → (12.4, False, False)
+      "13+"    → (13.0, True, False)
+      "14.6?"  → (14.6, False, True)
+      "7+"     → (7.0, True, False)
+      "耐"     → (-1.0, False, False)  # non-numeric
+      ""       → (0.0, False, False)
+    """
+    is_plus = "+" in level_str
+    is_ura = "?" in level_str
+    # Strip non-numeric except . + ?
+    clean = level_str.replace("+", "").replace("?", "").strip()
+    try:
+        val = float(clean) if clean else 0.0
+    except ValueError:
+        val = -1.0
+    return val, is_plus, is_ura
+
+
+def parse_maidata(content: str, song_id: str = "",
+                  maidata_path: str = "", audio_path: str = "") -> Song:
+    """
+    Parse a complete maidata.txt content string into a Song object.
+
+    Args:
+        content: The full text content of maidata.txt
+        song_id: Folder name / song identifier
+        maidata_path: Relative path to maidata.txt
+        audio_path: Relative path to track.mp3
+
+    Returns:
+        Song object with all parsed data
+    """
+    song = Song(song_id=song_id, maidata_path=maidata_path, audio_path=audio_path)
+
+    # Strip BOM if present
+    if content.startswith("\ufeff"):
+        content = content[1:]
+
+    lines = content.split("\n")
+    raw_notes: dict[int, str] = {}  # difficulty_index → raw note text
+    current_note_idx: Optional[int] = None
+    current_note_lines: list[str] = []
+
+    for line in lines:
+        line_stripped = line.strip()
+        if not line_stripped:
+            continue
+
+        if line_stripped.startswith("&"):
+            key, value = parse_header_line(line_stripped)
+
+            # If we were collecting note data, flush it
+            if current_note_idx is not None:
+                raw_notes[current_note_idx] = "\n".join(current_note_lines)
+                current_note_idx = None
+                current_note_lines = []
+
+            if key.startswith("inote_"):
+                # Note data block
+                idx_str = key.replace("inote_", "")
+                try:
+                    current_note_idx = int(idx_str)
+                except ValueError:
+                    current_note_idx = None
+                if value:
+                    current_note_lines.append(value)
+            elif key.startswith("lv_"):
+                idx_str = key.replace("lv_", "")
+                if not idx_str.isdigit():
+                    continue
+                idx = int(idx_str)
+                if not value:
+                    continue  # skip empty levels (no chart)
+                song.levels[idx] = value
+                level_val, is_plus, is_ura = parse_level_value(value)
+                if level_val > 0:
+                    # Create chart placeholder (notes parsed later)
+                    diff = Difficulty.from_index(idx)
+                    chart = Chart(
+                        difficulty_index=idx,
+                        difficulty=diff,
+                        level=value,
+                        level_value=level_val,
+                        is_plus=is_plus,
+                        is_ura=is_ura,
+                    )
+                    song.charts[idx] = chart
+            elif key.startswith("des_"):
+                idx_str = key.replace("des_", "")
+                try:
+                    idx = int(idx_str)
+                except ValueError:
+                    continue
+                song.charters[idx] = value
+                if idx in song.charts:
+                    song.charts[idx].charter = value
+            elif key == "title":
+                song.title = value
+                song.title_clean = _clean_title(value)
+                song.tags = _extract_tags(value)
+                song.is_full = "[FULL]" in value or "_FULLTOUCH" in value
+                song.is_fulltouch = "_FULLTOUCH" in value
+            elif key == "artist":
+                song.artist = value
+            elif key == "artistid":
+                try:
+                    song.artist_id = int(value)
+                except ValueError:
+                    pass
+            elif key == "wholebpm":
+                try:
+                    song.bpm = float(value)
+                except ValueError:
+                    pass
+            elif key == "first":
+                try:
+                    song.first = float(value)
+                except ValueError:
+                    pass
+            elif key == "genre":
+                song.genre = value
+                song.is_utage = "宴会場" in value
+            elif key == "genreid":
+                try:
+                    song.genre_id = int(value)
+                except ValueError:
+                    pass
+            elif key == "cabinet":
+                try:
+                    song.cabinet = Cabinet(value.upper())
+                except ValueError:
+                    song.cabinet = Cabinet.UNKNOWN
+            elif key == "version":
+                song.version = value
+            elif key == "shortid":
+                try:
+                    song.short_id = int(value)
+                except ValueError:
+                    pass
+            elif key == "des":
+                song.description = value
+            elif key == "ChartConverter":
+                song.converter = value
+            elif key == "ChartConvertTool":
+                song.converter_tool = value
+            elif key == "ChartConvertToolVersion":
+                song.converter_version = value
+        elif current_note_idx is not None:
+            # Collect note data lines
+            current_note_lines.append(line_stripped)
+
+    # Flush last note block
+    if current_note_idx is not None:
+        raw_notes[current_note_idx] = "\n".join(current_note_lines)
+
+    # ── Parse all collected note sequences ──
+    for idx, raw in raw_notes.items():
+        notes = _parse_note_sequence(raw)
+
+        if idx in song.charts:
+            song.charts[idx].notes = notes
+            song.charts[idx].compute_stats()
+        elif idx in song.levels and song.levels[idx]:
+            # Chart has notes + level but was skipped due to level_val being non-numeric (e.g. "耐")
+            level_val, is_plus, is_ura = parse_level_value(song.levels[idx])
+            diff = Difficulty.from_index(idx)
+            chart = Chart(
+                difficulty_index=idx,
+                difficulty=diff,
+                level=song.levels[idx],
+                level_value=level_val,
+                is_plus=is_plus,
+                is_ura=is_ura,
+            )
+            chart.notes = notes
+            chart.compute_stats()
+            song.charts[idx] = chart
+
+    return song
+
+
+def _clean_title(title: str) -> str:
+    """Remove bracket tags like [SD], [DX], [宴] from title."""
+    return re.sub(r"\[.*?\]", "", title).strip()
+
+
+def _extract_tags(title: str) -> list[str]:
+    """Extract bracket tags from title, e.g. [SD], [DX], [宴]."""
+    return re.findall(r"\[(.*?)\]", title)
+
+
+def parse_maidata_file(filepath: str | Path) -> Song:
+    """
+    Parse a maidata.txt file from disk.
+
+    Args:
+        filepath: Path to the maidata.txt file
+
+    Returns:
+        Parsed Song object
+    """
+    filepath = Path(filepath)
+    song_id = filepath.parent.name
+    content = filepath.read_text(encoding="utf-8")
+
+    # Build relative paths
+    maidata_rel = filepath.name
+    audio_rel = "track.mp3"
+
+    return parse_maidata(
+        content,
+        song_id=song_id,
+        maidata_path=maidata_rel,
+        audio_path=audio_rel,
+    )

model.py

@@ -0,0 +1,681 @@
+"""
+MaiGenerator — Audio-conditioned autoregressive maimai chart generator.
+
+Architecture: Encoder-Decoder Transformer with time-aligned RoPE.
+
+  Audio tokens → AudioEncoder → audio_feat [T_aud, d]
+                                     │ Cross-Attention
+  Chart tokens → ChartDecoder ───────┘
+  (autoregressive,  + BPM/Diff/Genre conditioning
+   causal mask)     + time-aligned RoPE positions
+
+Key design: Chart RoPE uses audio frame indices (via BPM translation),
+ensuring strong time alignment between music and generated notes.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import tokenizer as chart_tokenizer
+from tokenizer import CONFIG_BASE
+
+# ═══════════════════════════════════════════════════════════════════════
+# Constants (aligned with tokenizers)
+# ═══════════════════════════════════════════════════════════════════════
+
+CHART_VOCAB_SIZE = chart_tokenizer.VOCAB_SIZE
+AUDIO_VOCAB_SIZE = 2051      # MaiTrackTokenizer (2-layer EnCodec)
+AUDIO_FRAME_RATE = 75        # EnCodec 24kHz / 320 stride
+
+BOS, EOS, PAD = 1, 2, 0
+
+DIFF_NAMES = ["BASIC", "ADVANCED", "EXPERT", "MASTER", "ReMASTER"]
+NUM_DIFFICULTIES = len(DIFF_NAMES)
+
+# Beat division token → value
+DIV_MAP = {5: 1, 6: 2, 7: 4, 8: 8, 9: 16, 10: 32,
+           11: 48, 12: 64, 13: 128, 14: 192, 15: 384}
+DUR_TOKEN = 17  # [DUR] in chart vocab
+RST_TOKEN = 16
+TAP_BASE, TAP_END = 18, 26
+BRK_BASE, BRK_END = 26, 34
+HLD_BASE, HLD_END = 34, 42
+SLD_BASE, SLD_END = 42, 50
+SLD_BEG_TOKEN = 50
+SLD_END_TOKEN = 51
+SIM_BEG_TOKEN = 52
+SIM_END_TOKEN = 53
+TCH_BASE, TCH_END = 54, 95
+
+TYPE_REST = 0
+TYPE_TAP = 1
+TYPE_HOLD = 2
+TYPE_SLIDE = 3
+TYPE_BREAK = 4
+TYPE_TOUCH = 5
+TYPE_CONTROL = 6
+NUM_TOKEN_TYPES = 7
+NUM_POSITIONS = 9  # 0-7 real positions, 8 = none/control
+NUM_DIV_CLASSES = len(DIV_MAP)
+
+
+def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+          dropout_p: float = 0.0, is_causal: bool = False) -> torch.Tensor:
+    """Use PyTorch SDPA so CUDA can pick Flash/mem-efficient attention kernels."""
+    return F.scaled_dot_product_attention(
+        q, k, v,
+        dropout_p=dropout_p if torch.is_grad_enabled() else 0.0,
+        is_causal=is_causal,
+    )
+
+
+def is_timeline_token(tok: torch.Tensor) -> torch.Tensor:
+    """Tokens that correspond to one chart time slot after decoding."""
+    is_note = (((tok >= TAP_BASE) & (tok < TAP_END)) |
+               ((tok >= BRK_BASE) & (tok < BRK_END)) |
+               ((tok >= HLD_BASE) & (tok < HLD_END)) |
+               ((tok >= SLD_BASE) & (tok < SLD_END)) |
+               ((tok >= TCH_BASE) & (tok < TCH_END)) |
+               (tok >= CONFIG_BASE))
+    return (tok == RST_TOKEN) | (tok == SIM_BEG_TOKEN) | (tok == SLD_BEG_TOKEN) | is_note
+
+
+def token_structure_features(tokens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return token type ids and 0-based position ids for structure-aware embeds."""
+    typ = torch.full_like(tokens, TYPE_CONTROL)
+    pos = torch.full_like(tokens, NUM_POSITIONS - 1)
+
+    typ = torch.where(tokens == RST_TOKEN, torch.full_like(typ, TYPE_REST), typ)
+
+    ranges = [
+        (TAP_BASE, TAP_END, TYPE_TAP),
+        (HLD_BASE, HLD_END, TYPE_HOLD),
+        (SLD_BASE, SLD_END, TYPE_SLIDE),
+        (BRK_BASE, BRK_END, TYPE_BREAK),
+        (TCH_BASE, TCH_END, TYPE_TOUCH),
+    ]
+    for start, end, token_type in ranges:
+        mask = (tokens >= start) & (tokens < end)
+        typ = torch.where(mask, torch.full_like(typ, token_type), typ)
+        pos = torch.where(mask, (tokens - start) % 8, pos)
+    return typ, pos
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# RoPE with custom positions
+# ═══════════════════════════════════════════════════════════════════════
+
+class RoPE(nn.Module):
+    """Rotary Position Embedding supporting custom position indices."""
+
+    def __init__(self, dim: int, base: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, x: torch.Tensor,
+                positions: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: [B, H, T, D] — query or key.
+            positions: [T] or [B, T] — custom position indices.
+
+        Returns:
+            Rotated tensor [B, H, T, D].
+        """
+        B, H, T, D = x.shape
+        device = x.device
+
+        if positions is None:
+            positions = torch.arange(T, device=device, dtype=torch.float32)
+
+        if positions.dim() == 1:
+            positions = positions.unsqueeze(0)
+        positions = positions.unsqueeze(1).unsqueeze(-1)  # [B, 1, T, 1]
+
+        angles = positions * self.inv_freq.to(device)  # [B, 1, T, D/2]
+        sin, cos = angles.sin(), angles.cos()
+
+        x_even, x_odd = x[..., 0::2], x[..., 1::2]
+        out = torch.empty_like(x)
+        out[..., 0::2] = x_even * cos - x_odd * sin
+        out[..., 1::2] = x_even * sin + x_odd * cos
+        return out
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Onset Feature Injection (FiLM)
+# ═══════════════════════════════════════════════════════════════════════
+
+class OnsetFiLM(nn.Module):
+    """Beat-prior injection: modulates encoder features at beat positions."""
+
+    def __init__(self, d_model: int = 512):
+        super().__init__()
+        self.gamma = nn.Linear(1, d_model)
+        self.beta = nn.Linear(1, d_model)
+
+    def forward(self, enc_out, onset):
+        """enc_out: [B, T_enc, D], onset: [B, T_enc, 1]"""
+        g = torch.tanh(self.gamma(onset)) * 0.5  # [B, T_enc, D]
+        b = self.beta(onset) * 0.1
+        return enc_out * (1.0 + g) + b
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# MoE FFN (Mixture of Experts for difficulty routing)
+# ═══════════════════════════════════════════════════════════════════════
+
+class MoEFFN(nn.Module):
+    """MoE FFN: routes input through N experts weighted by difficulty."""
+
+    def __init__(self, d_model=512, d_ff=2048, n_experts=6, dropout=0.1):
+        super().__init__()
+        self.n_experts = n_experts
+        self.experts = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout),
+                nn.Linear(d_ff, d_model), nn.Dropout(dropout),
+            ) for _ in range(n_experts)
+        ])
+        self.router = nn.Linear(d_model, n_experts)
+
+    def forward(self, x, diff_emb):
+        """x: [B, T, d], diff_emb: [B, d]"""
+        weights = F.softmax(self.router(diff_emb), dim=-1)  # [B, N]
+        out = sum(weights[:, i:i+1, None] * self.experts[i](x)
+                  for i in range(self.n_experts))
+        return out
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Encoder Block
+# ═══════════════════════════════════════════════════════════════════════
+
+class EncoderBlock(nn.Module):
+    """Pre-LN encoder: Self-Attn + FFN. Uses FlashAttention for O(T) memory."""
+
+    def __init__(self, d_model: int = 512, heads: int = 8,
+                 d_ff: int = 2048, dropout: float = 0.1):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = d_model // heads
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.rope = RoPE(self.head_dim)
+
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(d_ff, d_model), nn.Dropout(dropout),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor,
+                positions: Optional[torch.Tensor] = None,
+                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, T, D = x.shape
+
+        residual = x
+        x_norm = self.norm1(x)
+
+        Q = self.q_proj(x_norm).view(B, T, self.heads, self.head_dim).transpose(1, 2)
+        K = self.k_proj(x_norm).view(B, T, self.heads, self.head_dim).transpose(1, 2)
+        V = self.v_proj(x_norm).view(B, T, self.heads, self.head_dim).transpose(1, 2)
+
+        # RoPE time-aligned positions.
+        if positions is not None:
+            Q = self.rope(Q, positions)
+            K = self.rope(K, positions)
+
+        if mask is not None:
+            attn_mask = mask
+            if attn_mask.dtype != torch.bool:
+                attn_mask = attn_mask.to(torch.bool)
+            attn_out = F.scaled_dot_product_attention(
+                Q, K, V,
+                attn_mask=attn_mask,
+                dropout_p=self.dropout.p if self.training else 0.0,
+            )
+        else:
+            attn_out = _sdpa(Q, K, V, self.dropout.p if self.training else 0.0)
+        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, D)
+        x = residual + self.dropout(self.out_proj(attn_out))
+
+        residual = x
+        x = residual + self.ffn(self.norm2(x))
+        return x
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Decoder Block
+# ═══════════════════════════════════════════════════════════════════════
+
+class DecoderBlock(nn.Module):
+    """Pre-LN decoder: Causal Self-Attn + Cross-Attn + FFN.
+    
+    Supports KV-cache for fast autoregressive inference.
+    """
+
+    def __init__(self, d_model: int = 512, heads: int = 8,
+                 d_ff: int = 2048, dropout: float = 0.1,
+                 use_moe: bool = False, n_experts: int = 6):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = d_model // heads
+
+        # Self-attn
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.rope = RoPE(self.head_dim)
+
+        # Cross-attn
+        self.cross_q = nn.Linear(d_model, d_model)
+        self.cross_k = nn.Linear(d_model, d_model)
+        self.cross_v = nn.Linear(d_model, d_model)
+        self.cross_out = nn.Linear(d_model, d_model)
+
+        # FFN (standard or MoE)
+        if use_moe:
+            self.ffn = MoEFFN(d_model, d_ff, n_experts, dropout)
+            self.is_moe = True
+        else:
+            self.ffn = nn.Sequential(
+                nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout),
+                nn.Linear(d_ff, d_model), nn.Dropout(dropout),
+            )
+            self.is_moe = False
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def _init_self_kv_cache(self, batch_size: int, max_len: int, device):
+        """Pre-allocate self-attention KV cache: tuple of (K, V) tensors."""
+        D = self.head_dim
+        H = self.heads
+        self._self_k_cache = torch.zeros(batch_size, H, max_len, D, device=device)
+        self._self_v_cache = torch.zeros(batch_size, H, max_len, D, device=device)
+        self._self_cache_len = 0
+
+    def _init_cross_kv_cache(self, enc_out: torch.Tensor):
+        """Precompute cross-attention K, V from encoder output (fixed during generation).
+        Note: enc_out is the raw encoder output, NOT pre-normed (norm2 is for decoder hidden)."""
+        B, T_enc, D = enc_out.shape
+        Kc = self.cross_k(enc_out).view(B, T_enc, self.heads, self.head_dim).transpose(1, 2)
+        Vc = self.cross_v(enc_out).view(B, T_enc, self.heads, self.head_dim).transpose(1, 2)
+        self._cross_k_cache = Kc
+        self._cross_v_cache = Vc
+
+    def forward(self, x, enc_out, self_positions=None, diff_emb=None,
+                onset_film=None, onset_kv=None,
+                use_cache: bool = False):
+        B, T_dec, D = x.shape
+
+        # ── Causal Self-Attn ──
+        residual = x
+        x_norm = self.norm1(x)
+        Q = self.q_proj(x_norm).view(B, T_dec, self.heads, self.head_dim).transpose(1, 2)
+        K = self.k_proj(x_norm).view(B, T_dec, self.heads, self.head_dim).transpose(1, 2)
+        V = self.v_proj(x_norm).view(B, T_dec, self.heads, self.head_dim).transpose(1, 2)
+
+        # RoPE time-aligned positions (full sequence positions for cache mode)
+        if self_positions is not None:
+            Q = self.rope(Q, self_positions)
+            K = self.rope(K, self_positions)
+
+        if use_cache and hasattr(self, '_self_k_cache'):
+            # Append new K, V to cache
+            cache_len = self._self_cache_len
+            new_len = T_dec
+            self._self_k_cache[:, :, cache_len:cache_len + new_len] = K
+            self._self_v_cache[:, :, cache_len:cache_len + new_len] = V
+            K = self._self_k_cache[:, :, :cache_len + new_len]
+            V = self._self_v_cache[:, :, :cache_len + new_len]
+            self._self_cache_len = cache_len + new_len
+
+        attn = _sdpa(Q, K, V, self.dropout.p if self.training else 0.0,
+                     is_causal=(not use_cache))
+        x = residual + self.dropout(self.out_proj(
+            attn.transpose(1, 2).contiguous().view(B, T_dec, D)))
+
+        # ── Cross-Attn (with optional OnsetFiLM modulation) ──
+        residual = x
+        x_norm = self.norm2(x)
+        Qc = self.cross_q(x_norm).view(B, T_dec, self.heads, self.head_dim).transpose(1, 2)
+
+        if use_cache and hasattr(self, '_cross_k_cache'):
+            Kc = self._cross_k_cache
+            Vc = self._cross_v_cache
+        else:
+            _enc = enc_out
+            if onset_film is not None and onset_kv is not None:
+                _enc = onset_film(_enc, onset_kv)
+            T_enc = _enc.shape[1]
+            Kc = self.cross_k(_enc).view(B, T_enc, self.heads, self.head_dim).transpose(1, 2)
+            Vc = self.cross_v(_enc).view(B, T_enc, self.heads, self.head_dim).transpose(1, 2)
+
+        attn_c = _sdpa(Qc, Kc, Vc, self.dropout.p if self.training else 0.0)
+        x = residual + self.dropout(self.cross_out(
+            attn_c.transpose(1, 2).contiguous().view(B, T_dec, D)))
+
+        # ── FFN ──
+        residual = x
+        if self.is_moe and diff_emb is not None:
+            x = residual + self.ffn(self.norm3(x), diff_emb)
+        else:
+            x = residual + self.ffn(self.norm3(x))
+        return x
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# MaiGenerator
+# ═══════════════════════════════════════════════════════════════════════
+
+class MaiGenerator(nn.Module):
+    """Audio-conditioned autoregressive maimai chart generator.
+
+    Input:
+        audio_tokens:  [B, T_aud]   — EnCodec tokens
+        chart_tokens:  [B, T_chart] — chart token sequence (train: full[:-1])
+        bpm:           [B, 1]       — BPM value
+        difficulty:    [B, 1]       — difficulty enum (0..4)
+        level_value:   [B, 1]       — numeric level (e.g. 12.4)
+        genre:         [B, 1]       — genre index (optional)
+
+    Output:
+        logits: [B, T_chart, chart_vocab] — next-token prediction
+    """
+
+    def __init__(self, d_model: int = 512, enc_layers: int = 6,
+                 dec_layers: int = 8, heads: int = 8, d_ff: int = 2048,
+                 dropout: float = 0.1, chart_vocab: int | None = None,
+                 audio_vocab: int = AUDIO_VOCAB_SIZE,
+                 num_genres: int = 16, max_audio_len: int = 32768,
+                 audio_downsample: int = 8, use_moe: bool = True,
+                 n_experts: int = 6, moe_layers: list = None):
+        """moe_layers: indices of decoder layers that use MoE (e.g. [8,9,10,11])
+        Other decoder layers use shared FFN. None = all decoder layers use MoE."""
+        super().__init__()
+        if chart_vocab is None:
+            chart_vocab = chart_tokenizer.VOCAB_SIZE
+        self.d_model = d_model
+        self.chart_vocab_size = chart_vocab
+        self.audio_downsample = audio_downsample
+
+        if moe_layers is None:
+            moe_layers = list(range(dec_layers))  # all MoE
+
+        # Embeddings
+        self.audio_embed = nn.Embedding(audio_vocab, d_model)
+        self.chart_embed = nn.Embedding(chart_vocab, d_model)
+        self.chart_type_embed = nn.Embedding(NUM_TOKEN_TYPES, d_model)
+        self.chart_pos_embed = nn.Embedding(NUM_POSITIONS, d_model)
+
+        # Onset feature injection
+        self.onset_film = OnsetFiLM(d_model)
+
+        # Audio downsampling (Conv1D + LayerNorm to reduce seq len)
+        if audio_downsample > 1:
+            self.audio_down = nn.Sequential(
+                nn.Conv1d(d_model, d_model, kernel_size=audio_downsample,
+                          stride=audio_downsample, padding=0),
+                nn.GELU(),
+            )
+        else:
+            self.audio_down = nn.Identity()
+
+        self.audio_pos_embed = nn.Embedding(max_audio_len, d_model)
+
+        # Conditions
+        self.bpm_proj = nn.Sequential(nn.Linear(1, d_model), nn.SiLU(),
+                                       nn.Linear(d_model, d_model))
+        self.diff_embed = nn.Embedding(NUM_DIFFICULTIES, d_model)
+        self.level_proj = nn.Sequential(nn.Linear(1, d_model), nn.SiLU(),
+                                         nn.Linear(d_model, d_model))
+        self.genre_embed = nn.Embedding(num_genres, d_model)
+
+        # Encoder / Decoder (hybrid: shared FFN + MoE layers)
+        self.audio_encoder = nn.ModuleList([
+            EncoderBlock(d_model, heads, d_ff, dropout) for _ in range(enc_layers)])
+        self.moe_layers = set(moe_layers)
+        n_shared = dec_layers - len(self.moe_layers)
+        n_moe = len(self.moe_layers)
+        print(f"Decoder: {n_shared} shared + {n_moe} MoE ×{n_experts} experts")
+        self.chart_decoder = nn.ModuleList([
+            DecoderBlock(d_model, heads, d_ff, dropout,
+                         use_moe=(i in self.moe_layers), n_experts=n_experts)
+            for i in range(dec_layers)])
+
+        self.output_head = nn.Linear(d_model, chart_vocab)
+        self.presence_head = nn.Linear(d_model, 2)
+        self.type_head = nn.Linear(d_model, NUM_TOKEN_TYPES)
+        self.position_head = nn.Linear(d_model, 8)
+        self.division_head = nn.Linear(d_model, NUM_DIV_CLASSES)
+        self.sim_head = nn.Linear(d_model, 2)
+        self.duration_head = nn.Linear(d_model, 2)
+        self.enc_norm = nn.LayerNorm(d_model)
+        self.dec_norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight, gain=0.5)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding):
+                nn.init.normal_(m.weight, std=0.02)
+
+    def init_kv_cache(self, batch_size: int, max_len: int, enc_out: torch.Tensor):
+        """Initialize KV caches for all decoder blocks. Call before incremental generation.
+        
+        Args:
+            batch_size: Batch size (usually 1 for inference).
+            max_len: Maximum generation length for pre-allocation.
+            enc_out: Encoder output [B, T_enc, D] for cross-attention cache.
+        """
+        for blk in self.chart_decoder:
+            blk._init_self_kv_cache(batch_size, max_len, enc_out.device)
+            blk._init_cross_kv_cache(enc_out)
+
+    # ── Time-aligned positions (核心) ──────────────────────────────
+
+    @staticmethod
+    def compute_chart_positions(chart_tokens: torch.Tensor,
+                                bpm: torch.Tensor,
+                                downsample: int = 4) -> torch.Tensor:
+        """Compute downsampled audio-frame positions for chart tokens.
+
+        Args:
+            chart_tokens: [B, T]
+            bpm: [B, 1]
+            downsample: Audio downsampling factor.
+
+        Returns:
+            positions: [B, T] — downsampled frame indices (float).
+        """
+        B, T = chart_tokens.shape
+        device = chart_tokens.device
+        bpm_v = bpm.view(B).float().clamp(min=30.0)
+
+        div_values = torch.full((B,), 4.0, device=device)
+        positions = torch.zeros(B, T, device=device)
+        current_beat = torch.zeros(B, device=device)
+        dur_param_skip = torch.zeros(B, dtype=torch.long, device=device)
+        sim_skip = torch.zeros(B, dtype=torch.long, device=device)
+        slide_active = torch.zeros(B, dtype=torch.bool, device=device)
+
+        for i in range(T):
+            tok = chart_tokens[:, i]
+
+            # Update beat division
+            for div_id, div_val in DIV_MAP.items():
+                div_values = torch.where(tok == div_id,
+                                         torch.tensor(float(div_val), device=device),
+                                         div_values)
+
+            beat_per_token = 4.0 / div_values
+
+            # Record position: beat → seconds → audio frame → downsampled
+            time_sec = current_beat * 60.0 / bpm_v
+            positions[:, i] = time_sec * AUDIO_FRAME_RATE / downsample
+
+            # Advance beat only for decoded timeline slots. SIM/SLD groups
+            # occupy one slot at their begin token; their contents are structural.
+            is_dur_param = dur_param_skip > 0
+            is_dur = (tok == DUR_TOKEN)
+            in_sim_body = sim_skip > 0
+            in_slide_body = slide_active & (tok != SLD_BEG_TOKEN)
+            group_body = in_sim_body | in_slide_body
+            advances_time = is_timeline_token(tok) & ~is_dur_param & ~is_dur & ~group_body
+            current_beat = torch.where(advances_time,
+                                       current_beat + beat_per_token,
+                                       current_beat)
+
+            is_sim_beg = tok == SIM_BEG_TOKEN
+            is_sim_end = tok == SIM_END_TOKEN
+            count_after_sim_beg = is_sim_beg & (i + 1 < T)
+            next_tok = chart_tokens[:, i + 1] if i + 1 < T else torch.zeros_like(tok)
+            sim_skip = torch.where(count_after_sim_beg,
+                                   torch.clamp(next_tok + 2, min=0),
+                                   torch.clamp(sim_skip - 1, min=0))
+            sim_skip = torch.where(is_sim_end, torch.zeros_like(sim_skip), sim_skip)
+
+            slide_active = torch.where(tok == SLD_BEG_TOKEN, torch.ones_like(slide_active), slide_active)
+            slide_active = torch.where(tok == SLD_END_TOKEN, torch.zeros_like(slide_active), slide_active)
+
+            dur_param_skip = torch.where(is_dur,
+                                         torch.full_like(dur_param_skip, 2),
+                                         torch.clamp(dur_param_skip - 1, min=0))
+
+        return positions
+
+    # ── Forward ────────────────────────────────────────────────────
+
+    def forward(self, audio_tokens: torch.Tensor,
+                chart_tokens: torch.Tensor, bpm: torch.Tensor,
+                difficulty: torch.Tensor, level_value: torch.Tensor,
+                genre: Optional[torch.Tensor] = None,
+                onset_curve: Optional[torch.Tensor] = None,
+                return_aux: bool = False,
+                return_hidden: bool = False) -> torch.Tensor | dict[str, torch.Tensor]:
+        """Training forward pass (teacher forcing).
+
+        Returns:
+            logits: [B, T_chart, 256]
+        """
+        B, T_chart = chart_tokens.shape
+        device = chart_tokens.device
+
+        # ── Encode audio ──
+        T_aud = audio_tokens.shape[1]
+        aud = self.audio_embed(audio_tokens)
+        aud = aud + self.audio_pos_embed(
+            torch.arange(T_aud, device=device).unsqueeze(0).expand(B, -1))
+
+        # Downsample: [B, T, D] → [B, T//stride, D]
+        if self.audio_downsample > 1:
+            aud = aud.transpose(1, 2)  # [B, D, T]
+            aud = self.audio_down(aud)  # [B, D, T']
+            aud = aud.transpose(1, 2)  # [B, T', D]
+            T_aud = aud.shape[1]
+
+        if onset_curve is None:
+            delta = torch.zeros(B, T_aud, device=device, dtype=aud.dtype)
+            if T_aud > 1:
+                delta[:, 1:] = (aud[:, 1:] - aud[:, :-1]).pow(2).mean(dim=-1).sqrt()
+                denom = delta.amax(dim=1, keepdim=True).clamp_min(1e-6)
+                delta = delta / denom
+            onset_curve = delta.unsqueeze(-1)
+        elif onset_curve.dim() == 2:
+            onset_curve = onset_curve.unsqueeze(-1)
+        if onset_curve.shape[1] != T_aud:
+            onset_curve = F.interpolate(onset_curve.transpose(1, 2), size=T_aud,
+                                        mode="linear", align_corners=False).transpose(1, 2)
+
+        aud = self.dropout(aud)
+        aud_pos = torch.arange(T_aud, device=device, dtype=torch.float32)
+        for blk in self.audio_encoder:
+            aud = blk(aud, positions=aud_pos)
+        aud = self.enc_norm(aud)
+
+        # Diff vector for decoder MoE routing
+        diff_vec = self.diff_embed(difficulty.squeeze(-1))  # [B, d_model]
+
+        # Embed chart + structural token features + conditions
+        token_type, token_pos = token_structure_features(chart_tokens)
+        emb = (self.chart_embed(chart_tokens) +
+               self.chart_type_embed(token_type) +
+               self.chart_pos_embed(token_pos))
+
+        bpm_emb = self.bpm_proj(bpm.float()).unsqueeze(1)
+        diff_emb = self.diff_embed(difficulty.squeeze(-1)).unsqueeze(1)
+        level_emb = self.level_proj(level_value.float()).unsqueeze(1)
+        genre_emb = torch.zeros(B, 1, self.d_model, device=device)
+        if genre is not None:
+            genre_emb = self.genre_embed(genre.squeeze(-1)).unsqueeze(1)
+
+        emb = emb + bpm_emb + diff_emb + level_emb + genre_emb
+
+        # ── Time-aligned positions ──
+        chart_pos = self.compute_chart_positions(chart_tokens, bpm,
+                                                  self.audio_downsample)
+
+        # ── Decode ──
+        x = emb
+        onset_kv = onset_curve if onset_curve is not None else None
+        for blk in self.chart_decoder:
+            x = blk(x, enc_out=aud, self_positions=chart_pos,
+                     diff_emb=diff_vec, onset_film=self.onset_film, onset_kv=onset_kv)
+        x = self.dec_norm(x)
+
+        if return_hidden:
+            result = {"hidden": x}
+            if return_aux:
+                result.update({
+                    "presence": self.presence_head(x),
+                    "type": self.type_head(x),
+                    "position": self.position_head(x),
+                    "division": self.division_head(x),
+                    "sim": self.sim_head(x),
+                    "duration": self.duration_head(x),
+                })
+            return result
+
+        token_logits = self.output_head(x)
+        if not return_aux:
+            return token_logits
+
+        return {
+            "token": token_logits,
+            "presence": self.presence_head(x),
+            "type": self.type_head(x),
+            "position": self.position_head(x),
+            "division": self.division_head(x),
+            "sim": self.sim_head(x),
+            "duration": self.duration_head(x),
+        }
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device

targets.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import torch
+
+from model import DIV_MAP
+from tokenizer import (
+    PAD, RST, SIM_BEG, SIM_END, SLD_BEG, DUR, SIM_COUNT_2,
+    BRK_BASE, BRK_END, HLD_BASE, HLD_END, ID_TO_DIV,
+    DUR_NUM_TO_ID, DUR_DEN_TO_ID,
+    SLD_BASE, SLD_END, TAP_BASE, TAP_END, TCH_BASE, TCH_END,
+    CONFIG_BASE,
+)
+
+TYPE_REST = 0
+TYPE_TAP = 1
+TYPE_HOLD = 2
+TYPE_SLIDE = 3
+TYPE_BREAK = 4
+TYPE_TOUCH = 5
+TYPE_CONTROL = 6
+NUM_TOKEN_TYPES = 7
+
+DIV_VALUES = [1, 2, 4, 8, 16, 32, 48, 64, 128, 192, 384]
+DIV_VALUE_TO_CLASS = {v: i for i, v in enumerate(DIV_VALUES)}
+NUM_DIV_CLASSES = len(DIV_VALUES)
+
+
+def build_aux_targets(tokens: torch.Tensor) -> dict[str, torch.Tensor]:
+    """Build auxiliary labels from next-token targets.
+
+    Args:
+        tokens: [B, T] target token ids, usually chart_tokens[:, 1:].
+    """
+    device = tokens.device
+    B, T = tokens.shape
+
+    presence = torch.zeros(B, T, dtype=torch.long, device=device)
+    token_type = torch.full((B, T), TYPE_CONTROL, dtype=torch.long, device=device)
+    position = torch.full((B, T), -100, dtype=torch.long, device=device)
+    division = torch.full((B, T), -100, dtype=torch.long, device=device)
+    is_sim = torch.zeros(B, T, dtype=torch.long, device=device)
+    needs_duration = torch.zeros(B, T, dtype=torch.long, device=device)
+
+    is_pad = tokens == PAD
+    token_type = torch.where(tokens == RST, torch.full_like(token_type, TYPE_REST), token_type)
+
+    ranges = [
+        (TAP_BASE, TAP_END, TYPE_TAP),
+        (HLD_BASE, HLD_END, TYPE_HOLD),
+        (SLD_BASE, SLD_END, TYPE_SLIDE),
+        (BRK_BASE, BRK_END, TYPE_BREAK),
+        (TCH_BASE, TCH_END, TYPE_TOUCH),
+    ]
+    for start, end, typ in ranges:
+        mask = (tokens >= start) & (tokens < end)
+        token_type = torch.where(mask, torch.full_like(token_type, typ), token_type)
+        presence = torch.where(mask, torch.ones_like(presence), presence)
+        position = torch.where(mask, (tokens - start) % 8, position)
+
+    group_event = (tokens == SIM_BEG) | (tokens == SLD_BEG)
+    presence = torch.where(group_event, torch.ones_like(presence), presence)
+    config_event = tokens >= CONFIG_BASE
+    presence = torch.where(config_event, torch.ones_like(presence), presence)
+    sim_event = (tokens == SIM_BEG) | (tokens == SIM_COUNT_2) | (tokens == SIM_END)
+    is_sim = torch.where(sim_event, torch.ones_like(is_sim), is_sim)
+    duration_tokens = {DUR} | set(DUR_NUM_TO_ID.values()) | set(DUR_DEN_TO_ID.values())
+    duration_event = torch.zeros_like(needs_duration, dtype=torch.bool)
+    for tok in duration_tokens:
+        duration_event |= tokens == tok
+    needs_duration = torch.where(duration_event, torch.ones_like(needs_duration), needs_duration)
+
+    for div_id, div_value in ID_TO_DIV.items():
+        mask = tokens == div_id
+        division = torch.where(mask, torch.full_like(division, DIV_VALUE_TO_CLASS[div_value]), division)
+
+    token_type = torch.where(is_pad, torch.full_like(token_type, -100), token_type)
+    presence = torch.where(is_pad, torch.full_like(presence, -100), presence)
+    is_sim = torch.where(is_pad, torch.full_like(is_sim, -100), is_sim)
+    needs_duration = torch.where(is_pad, torch.full_like(needs_duration, -100), needs_duration)
+
+    return {
+        "presence": presence,
+        "type": token_type,
+        "position": position,
+        "division": division,
+        "sim": is_sim,
+        "duration": needs_duration,
+    }

tokenizer.py

@@ -0,0 +1,1145 @@
+"""
+Maimai Chart Tokenizer — rule-based bidirectional chart ↔ token conversion.
+
+Design:
+  - BPM is NOT tokenized (computed separately by external program)
+  - Beat division (div_N) tokens control note granularity
+  - Each note event → 1~5 tokens, lossless round-trip
+
+Vocabulary size: 256 (0-255), with room for expansion.
+
+Usage:
+    from tokenizer import MaiChartTokenizer
+
+    tok = MaiChartTokenizer()
+    tokens = tok.encode(chart)
+    chart2 = tok.decode(tokens)  # lossless
+"""
+
+from __future__ import annotations
+from dataclasses import dataclass
+import json
+from typing import Optional
+
+from mai_parser.models import Chart, TouchNote
+
+# ═══════════════════════════════════════════════════════════════════════
+# Vocabulary definition
+# ═══════════════════════════════════════════════════════════════════════
+
+# --- Special tokens (0-4) ---
+PAD = 0
+BOS = 1
+EOS = 2
+SEP = 3
+MASK = 4
+_SPECIAL_END = 5
+
+# --- Beat division tokens (5-16): {1} {2} {4} {8} {16} {32} {48} {64} {128} {192} {384} ---
+_DIV_VALUES = [1, 2, 4, 8, 16, 32, 48, 64, 128, 192, 384]
+DIV_BASE = _SPECIAL_END
+DIV_TO_ID: dict[int, int] = {v: DIV_BASE + i for i, v in enumerate(_DIV_VALUES)}
+ID_TO_DIV: dict[int, int] = {v: k for k, v in DIV_TO_ID.items()}
+DIV_END = DIV_BASE + len(_DIV_VALUES)
+
+# --- Rest token (17) ---
+RST = DIV_END
+_RST_END = RST + 1
+
+# --- Duration marker (18): followed by 2 tokens [beat, subdiv] ---
+DUR = _RST_END
+_DUR_END = DUR + 1
+
+# --- Tap tokens (19-26): tap_1 ~ tap_8 ---
+TAP_BASE = _DUR_END
+TAP_TO_ID = {i: TAP_BASE + i - 1 for i in range(1, 9)}
+ID_TO_TAP = {v: k for k, v in TAP_TO_ID.items()}
+TAP_END = TAP_BASE + 8
+
+# --- Break tokens (27-34): brk_1 ~ brk_8 ---
+BRK_BASE = TAP_END
+BRK_TO_ID = {i: BRK_BASE + i - 1 for i in range(1, 9)}
+ID_TO_BRK = {v: k for k, v in BRK_TO_ID.items()}
+BRK_END = BRK_BASE + 8
+
+# --- Hold tokens (35-42): hld_1 ~ hld_8 ---
+HLD_BASE = BRK_END
+HLD_TO_ID = {i: HLD_BASE + i - 1 for i in range(1, 9)}
+ID_TO_HLD = {v: k for k, v in HLD_TO_ID.items()}
+HLD_END = HLD_BASE + 8
+
+# --- Slide waypoint tokens (43-50): sld_1 ~ sld_8 ---
+SLD_BASE = HLD_END
+SLD_TO_ID = {i: SLD_BASE + i - 1 for i in range(1, 9)}
+ID_TO_SLD = {v: k for k, v in SLD_TO_ID.items()}
+SLD_END = SLD_BASE + 8
+
+# --- Slide control (51-52) ---
+SLD_BEG = SLD_END        # slide start, next token = point count
+SLD_END_TOKEN = SLD_BEG + 1  # slide end marker
+_SLD_CTRL_END = SLD_END_TOKEN + 1
+
+# --- Simultaneous control (53-54) ---
+SIM_BEG = _SLD_CTRL_END  # simultaneous start, next token = note count
+SIM_END = SIM_BEG + 1     # simultaneous end
+_SIM_CTRL_END = SIM_END + 1
+
+# --- Touch tokens (55-90): A1-A8, B1-B8, C1-C8, D1-D8, E1-E8 ---
+TCH_BASE = _SIM_CTRL_END
+_TOUCH_ZONES = ["A", "B", "C", "D", "E"]
+
+_tch_map: dict[str, int] = {}
+_idx = TCH_BASE
+for zone in _TOUCH_ZONES:
+    for pos in range(1, 9):
+        _tch_map[f"{zone}{pos}"] = _idx
+        _idx += 1
+# C (center touch without position)
+_tch_map["C"] = _idx
+_idx += 1
+
+TCH_TO_ID = _tch_map
+ID_TO_TCH = {v: k for k, v in TCH_TO_ID.items()}
+TCH_END = _idx
+
+# --- Parameter tokens (96-120): context-safe count/duration values ---
+SIM_COUNT_2 = TCH_END
+_SIM_COUNT_END = SIM_COUNT_2 + 1
+
+_DUR_NUM_VALUES = [1, 2, 3, 4, 6, 8, 12, 16]
+DUR_NUM_BASE = _SIM_COUNT_END
+DUR_NUM_TO_ID = {v: DUR_NUM_BASE + i for i, v in enumerate(_DUR_NUM_VALUES)}
+ID_TO_DUR_NUM = {v: k for k, v in DUR_NUM_TO_ID.items()}
+DUR_NUM_END = DUR_NUM_BASE + len(_DUR_NUM_VALUES)
+
+_DUR_DEN_VALUES = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64]
+DUR_DEN_BASE = DUR_NUM_END
+DUR_DEN_TO_ID = {v: DUR_DEN_BASE + i for i, v in enumerate(_DUR_DEN_VALUES)}
+ID_TO_DUR_DEN = {v: k for k, v in DUR_DEN_TO_ID.items()}
+DUR_DEN_END = DUR_DEN_BASE + len(_DUR_DEN_VALUES)
+
+# --- Metadata/header/helper tokens (220-231): kept for backward compatibility ---
+META_BPM = 220      # followed by BPM//2 as uint8
+META_DIFF = 221     # followed by difficulty enum (0-4)
+META_LEVEL = 222    # followed by level*10 as uint8
+META_GENRE = 223    # followed by genre id
+META_END = 224      # end of metadata header
+SLD_MID = 229       # slide intermediate waypoint
+HLD_ON = 230        # hold is ongoing at this beat position (informational)
+SLD_ON = 231        # slide is ongoing at this beat position (informational)
+
+# --- Whole time-slot configuration tokens (256+) ---
+CONFIG_BASE = 256
+
+
+def _duration_pairs(max_beats: float = 4.0) -> list[tuple[int, int]]:
+    pairs: list[tuple[int, int]] = []
+    for n in _DUR_NUM_VALUES:
+        for d in _DUR_DEN_VALUES:
+            if n / d <= max_beats + 1e-9:
+                pairs.append((n, d))
+    return pairs
+
+
+CONFIG_DURATIONS = _duration_pairs(4.0)
+CONFIG_TO_ID: dict[tuple, int] = {}
+ID_TO_CONFIG: dict[int, tuple] = {}
+
+
+def _normalize_config_spec(spec) -> tuple:
+    return tuple(tuple(x) if isinstance(x, list) else x for x in spec)
+
+
+def _add_config(spec: tuple) -> int:
+    spec = _normalize_config_spec(spec)
+    if spec in CONFIG_TO_ID:
+        return CONFIG_TO_ID[spec]
+    idx = CONFIG_BASE + len(CONFIG_TO_ID)
+    CONFIG_TO_ID[spec] = idx
+    ID_TO_CONFIG[idx] = spec
+    if "_TOKEN_NAMES" in globals():
+        _TOKEN_NAMES[idx] = "cfg_" + "_".join(str(x) for x in spec)
+    global VOCAB_SIZE
+    VOCAB_SIZE = CONFIG_BASE + len(CONFIG_TO_ID)
+    if "MaiChartTokenizer" in globals():
+        MaiChartTokenizer.vocab_size = VOCAB_SIZE
+    return idx
+
+
+def _build_config_vocab() -> None:
+    # Single button/touch events.
+    for pos in range(1, 9):
+        _add_config(("tap", pos))
+        _add_config(("brk", pos))
+        for dur in CONFIG_DURATIONS:
+            _add_config(("hld", pos, dur[0], dur[1]))
+
+    for region in sorted(TCH_TO_ID):
+        _add_config(("tch", region))
+
+    # Two-note simultaneous button configurations. Holds share one duration.
+    button_types = ("tap", "brk", "hld")
+    for p1 in range(1, 9):
+        for p2 in range(p1 + 1, 9):
+            for t1 in button_types:
+                for t2 in button_types:
+                    if (t1, p1) > (t2, p2):
+                        continue
+                    if "hld" in (t1, t2):
+                        for dur in CONFIG_DURATIONS:
+                            _add_config(("pair", t1, p1, t2, p2, dur[0], dur[1]))
+                    else:
+                        _add_config(("pair", t1, p1, t2, p2))
+
+    # Common slide configurations: 2- and 3-point paths with duration.
+    for a in range(1, 9):
+        for b in range(1, 9):
+            if b == a:
+                continue
+            for dur in CONFIG_DURATIONS:
+                _add_config(("sld", a, b, dur[0], dur[1]))
+            for c in range(1, 9):
+                if c in (a, b):
+                    continue
+                for dur in CONFIG_DURATIONS:
+                    _add_config(("sld", a, b, c, dur[0], dur[1]))
+
+
+_build_config_vocab()
+
+VOCAB_SIZE = CONFIG_BASE + len(CONFIG_TO_ID)
+TOKENIZER_VERSION = 3
+
+
+def export_config_vocab() -> list[list]:
+    return [list(spec) for spec, _ in sorted(CONFIG_TO_ID.items(), key=lambda x: x[1])]
+
+
+def load_config_vocab(specs: list) -> None:
+    for spec in specs:
+        _add_config(tuple(spec))
+
+
+def save_config_vocab(path: str) -> None:
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(export_config_vocab(), f, ensure_ascii=False)
+
+
+def load_config_vocab_file(path: str) -> None:
+    with open(path, "r", encoding="utf-8") as f:
+        load_config_vocab(json.load(f))
+
+# --- Special tokens that start a multi-token group ---
+_MULTI_TOKEN_STARTS = {DUR, SLD_BEG, SLD_END_TOKEN, SIM_BEG, SIM_END}
+
+# ═══════════════════════════════════════════════════════════════════════
+# Token name lookup (for debugging)
+# ═══════════════════════════════════════════════════════════════════════
+
+_TOKEN_NAMES: dict[int, str] = {
+    PAD: "[PAD]",
+    BOS: "[BOS]",
+    EOS: "[EOS]",
+    SEP: "[SEP]",
+    MASK: "[MASK]",
+    RST: "[RST]",
+    DUR: "[DUR]",
+    SLD_BEG: "[SLD_BEG]",
+    SLD_END_TOKEN: "[SLD_END]",
+    SIM_BEG: "[SIM_BEG]",
+    SIM_END: "[SIM_END]",
+    META_BPM: "[META_BPM]",
+    META_DIFF: "[META_DIFF]",
+    META_LEVEL: "[META_LEVEL]",
+    META_GENRE: "[META_GENRE]",
+    META_END: "[META_END]",
+    SLD_MID: "[SLD_MID]",
+    HLD_ON: "[HLD_ON]",
+    SLD_ON: "[SLD_ON]",
+}
+for v, i in DIV_TO_ID.items():
+    _TOKEN_NAMES[i] = f"div_{v}"
+for p, i in TAP_TO_ID.items():
+    _TOKEN_NAMES[i] = f"tap_{p}"
+for p, i in BRK_TO_ID.items():
+    _TOKEN_NAMES[i] = f"brk_{p}"
+for p, i in HLD_TO_ID.items():
+    _TOKEN_NAMES[i] = f"hld_{p}"
+for p, i in SLD_TO_ID.items():
+    _TOKEN_NAMES[i] = f"sld_{p}"
+for t, i in TCH_TO_ID.items():
+    _TOKEN_NAMES[i] = f"tch_{t}"
+_TOKEN_NAMES[SIM_COUNT_2] = "sim_count_2"
+for v, i in DUR_NUM_TO_ID.items():
+    _TOKEN_NAMES[i] = f"dur_num_{v}"
+for v, i in DUR_DEN_TO_ID.items():
+    _TOKEN_NAMES[i] = f"dur_den_{v}"
+for spec, i in CONFIG_TO_ID.items():
+    _TOKEN_NAMES[i] = "cfg_" + "_".join(str(x) for x in spec)
+
+
+def token_name(token_id: int) -> str:
+    """Human-readable name for a token ID."""
+    return _TOKEN_NAMES.get(token_id, f"<{token_id}>")
+
+
+def _nearest(values: list[int], value: int) -> int:
+    return min(values, key=lambda x: abs(x - value))
+
+
+def encode_duration_tokens(duration: tuple[int, int]) -> list[int]:
+    beat = _nearest(_DUR_NUM_VALUES, max(1, int(duration[0])))
+    den = _nearest(_DUR_DEN_VALUES, max(1, int(duration[1])))
+    return [DUR, DUR_NUM_TO_ID[beat], DUR_DEN_TO_ID[den]]
+
+
+def read_duration_tokens(tokens: list[int], start: int) -> Optional[tuple[int, int]]:
+    if start + 2 >= len(tokens) or tokens[start] != DUR:
+        return None
+    num_tok = tokens[start + 1]
+    den_tok = tokens[start + 2]
+    if num_tok in ID_TO_DUR_NUM and den_tok in ID_TO_DUR_DEN:
+        return ID_TO_DUR_NUM[num_tok], ID_TO_DUR_DEN[den_tok]
+
+    # Backward compatibility with old checkpoints/caches that used raw ints.
+    beat = _nearest(_DUR_NUM_VALUES, max(1, min(int(num_tok), 16)))
+    den = _nearest(_DUR_DEN_VALUES, max(1, int(den_tok)))
+    return beat, den
+
+
+def make_sim_tokens(note_tokens: list[int]) -> list[int]:
+    note_tokens = [t for t in note_tokens if t not in (PAD, BOS, EOS)]
+    if len(note_tokens) <= 1:
+        return note_tokens
+    result: list[int] = []
+    result.extend([SIM_BEG, SIM_COUNT_2, note_tokens[0], note_tokens[1], SIM_END])
+    result.extend(note_tokens[2:])
+    return result
+
+
+def _snap_config_duration(duration: tuple[int, int] | None) -> tuple[int, int]:
+    if not duration:
+        return (1, 1)
+    n = _nearest(_DUR_NUM_VALUES, max(1, int(duration[0])))
+    d = _nearest(_DUR_DEN_VALUES, max(1, int(duration[1])))
+    if n / d > 4.0:
+        return min(CONFIG_DURATIONS, key=lambda x: (abs((x[0] / x[1]) - 4.0), x[1]))
+    return n, d
+
+
+def config_token_for_note(note: TouchNote) -> int | None:
+    if note.is_rest or note.is_end:
+        return None
+    if note.is_touch and len(note.touch_regions) == 1 and not note.is_hold:
+        return CONFIG_TO_ID.get(("tch", note.touch_regions[0]))
+    if note.is_slide:
+        path = note.slide_path or note.positions
+        if len(path) >= 2:
+            dur = _snap_config_duration(note.hold_duration)
+            spec = ("sld", *path, dur[0], dur[1])
+            if spec in CONFIG_TO_ID:
+                return CONFIG_TO_ID[spec]
+        return None
+
+    if len(note.positions) == 1:
+        pos = note.positions[0]
+        if not (1 <= pos <= 8):
+            return None
+        if note.is_hold:
+            dur = _snap_config_duration(note.hold_duration)
+            return CONFIG_TO_ID.get(("hld", pos, dur[0], dur[1]))
+        if note.is_break:
+            return CONFIG_TO_ID.get(("brk", pos))
+        return CONFIG_TO_ID.get(("tap", pos))
+
+    if len(note.positions) == 2:
+        p1, p2 = sorted(note.positions)
+        if not (1 <= p1 <= 8 and 1 <= p2 <= 8):
+            return None
+        typ = "hld" if note.is_hold else "brk" if note.is_break else "tap"
+        if typ == "hld":
+            dur = _snap_config_duration(note.hold_duration)
+            return CONFIG_TO_ID.get(("pair", "hld", p1, "hld", p2, dur[0], dur[1]))
+        return CONFIG_TO_ID.get(("pair", typ, p1, typ, p2))
+
+    return None
+
+
+def learn_config_from_note(note: TouchNote) -> int | None:
+    """Register a config token from a real chart note, preserving rare shapes."""
+    if note.is_rest or note.is_end:
+        return None
+    existing = config_token_for_note(note)
+    if existing is not None:
+        return existing
+
+    if note.is_touch and note.touch_regions and not note.is_hold:
+        return _add_config(("touch_multi", *sorted(note.touch_regions)))
+
+    if note.is_slide:
+        path = note.slide_path or note.positions
+        if len(path) >= 2:
+            dur = _snap_config_duration(note.hold_duration)
+            return _add_config(("sld", *path, dur[0], dur[1]))
+        return None
+
+    if len(note.positions) >= 2:
+        positions = sorted(p for p in note.positions if 1 <= p <= 8)
+        if len(positions) < 2:
+            return None
+        typ = "hld" if note.is_hold else "brk" if note.is_break else "tap"
+        if note.is_hold:
+            dur = _snap_config_duration(note.hold_duration)
+            spec = ("multi", typ, *positions, dur[0], dur[1])
+        else:
+            spec = ("multi", typ, *positions)
+        if spec in CONFIG_TO_ID:
+            return CONFIG_TO_ID[spec]
+        if len(positions) == 2:
+            return config_token_for_note(note)
+        return _add_config(spec)
+
+    return None
+
+
+def learn_config_vocab_from_charts(charts) -> int:
+    before = len(CONFIG_TO_ID)
+    for chart in charts:
+        for note in chart.notes:
+            learn_config_from_note(note)
+    return len(CONFIG_TO_ID) - before
+
+
+def note_from_config_token(token_id: int, beat_div: int) -> TouchNote | None:
+    spec = ID_TO_CONFIG.get(token_id)
+    if spec is None:
+        return None
+    note = TouchNote(beat_div=beat_div)
+    kind = spec[0]
+    if kind == "tap":
+        note.positions = [int(spec[1])]
+    elif kind == "brk":
+        note.positions = [int(spec[1])]
+        note.is_break = True
+    elif kind == "hld":
+        note.positions = [int(spec[1])]
+        note.is_hold = True
+        note.hold_duration = (int(spec[2]), int(spec[3]))
+    elif kind == "tch":
+        note.is_touch = True
+        note.touch_regions = [str(spec[1])]
+    elif kind == "touch_multi":
+        note.is_touch = True
+        note.touch_regions = [str(x) for x in spec[1:]]
+        note.is_simultaneous = len(note.touch_regions) > 1
+    elif kind == "pair":
+        t1, p1, t2, p2 = spec[1], int(spec[2]), spec[3], int(spec[4])
+        note.positions = [p1, p2]
+        note.is_simultaneous = True
+        if t1 == "hld" or t2 == "hld":
+            note.is_hold = True
+            note.hold_duration = (int(spec[5]), int(spec[6]))
+        elif t1 == "brk" or t2 == "brk":
+            note.is_break = True
+    elif kind == "sld":
+        *path, n, d = spec[1:]
+        note.positions = [int(x) for x in path]
+        note.slide_path = list(note.positions)
+        note.is_slide = True
+        note.hold_duration = (int(n), int(d))
+    elif kind == "multi":
+        typ = spec[1]
+        if typ == "hld":
+            *positions, n, d = spec[2:]
+            note.positions = [int(x) for x in positions]
+            note.is_hold = True
+            note.hold_duration = (int(n), int(d))
+        else:
+            note.positions = [int(x) for x in spec[2:]]
+            note.is_break = typ == "brk"
+        note.is_simultaneous = len(note.positions) > 1
+    else:
+        return None
+    return note
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Tokenizer class
+# ═══════════════════════════════════════════════════════════════════════
+
+class MaiChartTokenizer:
+    """
+    Rule-based bidirectional tokenizer for maimai charts.
+
+    encode(chart) → list[int]       # chart → tokens
+    decode(tokens) → Chart           # tokens → chart (lossless)
+    """
+
+    vocab_size: int = VOCAB_SIZE
+    pad_token_id: int = PAD
+    bos_token_id: int = BOS
+    eos_token_id: int = EOS
+    mask_token_id: int = MASK
+
+    # ── Encode ──────────────────────────────────────────────────────
+
+    def encode(self, chart: Chart, add_bos: bool = True,
+               add_eos: bool = True) -> list[int]:
+        """
+        Convert a Chart's notes into a token sequence.
+
+        Args:
+            chart: Parsed Chart from mai_parser.
+            add_bos: Prepend [BOS] token.
+            add_eos: Append [EOS] token.
+
+        Returns:
+            List of token IDs.
+        """
+        tokens: list[int] = []
+        if add_bos:
+            tokens.append(BOS)
+
+        current_div = 4  # default beat division
+
+        for note in chart.notes:
+            # Update beat division if changed
+            if note.beat_div != current_div:
+                current_div = note.beat_div
+                div_id = DIV_TO_ID.get(current_div)
+                if div_id is not None:
+                    tokens.append(div_id)
+
+            # Encode the note
+            tokens.extend(self._encode_note(note))
+
+        if add_eos:
+            tokens.append(EOS)
+
+        return tokens
+
+    def _encode_note(self, note: TouchNote) -> list[int]:
+        """Encode a single TouchNote → list of token IDs."""
+
+        if note.is_end:
+            return [EOS]
+
+        if note.is_rest:
+            return [RST]
+
+        cfg = config_token_for_note(note)
+        if cfg is not None:
+            return [cfg]
+
+        # Touch note
+        if note.is_touch:
+            return self._encode_touch(note)
+
+        # Break note
+        if note.is_break:
+            result = []
+            for pos in note.positions:
+                if 1 <= pos <= 8:
+                    result.append(BRK_TO_ID[pos])
+            return make_sim_tokens(result)
+
+        # Hold note
+        if note.is_hold:
+            result = []
+            for pos in note.positions:
+                if 1 <= pos <= 8:
+                    result.append(HLD_TO_ID[pos])
+            result = make_sim_tokens(result)
+            if note.hold_duration:
+                result.extend(encode_duration_tokens(note.hold_duration))
+            return result
+
+        # Slide note
+        if note.is_slide:
+            result = []
+            # Collect all positions (slide path)
+            positions = list(note.positions)
+            if note.slide_path:
+                # Use slide_path if available (more accurate)
+                positions = note.slide_path
+
+            if len(positions) >= 2:
+                result.append(SLD_BEG)
+                result.append(len(positions))
+                for pos in positions:
+                    if 1 <= pos <= 8:
+                        result.append(SLD_TO_ID[pos])
+                result.append(SLD_END_TOKEN)
+            elif len(positions) == 1 and 1 <= positions[0] <= 8:
+                result.append(SLD_TO_ID[positions[0]])
+
+            if note.hold_duration:
+                result.extend(encode_duration_tokens(note.hold_duration))
+            return result
+
+        # Regular tap
+        if len(note.positions) > 1:
+            result = []
+            for pos in note.positions:
+                if 1 <= pos <= 8:
+                    result.append(TAP_TO_ID[pos])
+            return make_sim_tokens(result)
+
+        # Single tap
+        for pos in note.positions:
+            if 1 <= pos <= 8:
+                return [TAP_TO_ID[pos]]
+
+        return [RST]  # fallback
+
+    def _encode_touch(self, note: TouchNote) -> list[int]:
+        """Encode a touch note."""
+        result = []
+        for region in note.touch_regions:
+            tid = TCH_TO_ID.get(region)
+            if tid is not None:
+                result.append(tid)
+        if len(result) > 1:
+            result = make_sim_tokens(result)
+        if note.is_hold and note.hold_duration:
+            result.extend(encode_duration_tokens(note.hold_duration))
+        return result if result else [RST]
+
+    # ── Decode ──────────────────────────────────────────────────────
+
+    def decode(self, tokens: list[int]) -> Chart:
+        """
+        Convert a token sequence back into a Chart.
+
+        Args:
+            tokens: List of token IDs (may include BOS/EOS).
+
+        Returns:
+            Reconstructed Chart (notes only; metadata not recoverable from tokens alone).
+        """
+        notes: list[TouchNote] = []
+        current_div = 4
+        i = 0
+
+        while i < len(tokens):
+            tid = tokens[i]
+
+            # Skip BOS
+            if tid == BOS:
+                i += 1
+                continue
+
+            # End of sequence
+            if tid == EOS:
+                note = TouchNote(beat_div=current_div, raw="E")
+                note.is_end = True
+                notes.append(note)
+                i += 1
+                continue
+
+            # Beat division change
+            if tid in ID_TO_DIV:
+                current_div = ID_TO_DIV[tid]
+                i += 1
+                continue
+
+            # Rest
+            if tid == RST:
+                note = TouchNote(beat_div=current_div, raw="")
+                note.is_rest = True
+                notes.append(note)
+                i += 1
+                continue
+
+            cfg_note = note_from_config_token(tid, current_div)
+            if cfg_note is not None:
+                notes.append(cfg_note)
+                i += 1
+                continue
+
+            # Duration marker → read next 2 tokens as (beat, subdiv)
+            # (handled inline in note decoding below)
+
+            # Slide start
+            if tid == SLD_BEG:
+                i += 1
+                if i >= len(tokens):
+                    break
+                if tokens[i] in ID_TO_SLD or tokens[i] in (SLD_MID, SLD_ON):
+                    # Unfolded slide format used by training/inference:
+                    # SLD_BEG [SLD_ON] sld_a [SLD_MID/SLD_ON sld_b ...] SLD_END [DUR beat subdiv]
+                    positions = []
+                    while i < len(tokens) and tokens[i] not in (SLD_END_TOKEN, DUR, EOS):
+                        if tokens[i] in (SLD_MID, SLD_ON):
+                            i += 1
+                            continue
+                        pt = tokens[i]
+                        if pt in ID_TO_SLD:
+                            positions.append(ID_TO_SLD[pt])
+                            i += 1
+                            continue
+                        # Malformed slide: stop before consuming unrelated chart events.
+                        break
+                    if i < len(tokens) and tokens[i] == SLD_END_TOKEN:
+                        i += 1
+                    dur = self._read_dur(tokens, i)
+                    if dur:
+                        i += 3
+                    note = TouchNote(beat_div=current_div, positions=positions)
+                    note.is_slide = True
+                    note.slide_path = list(positions)
+                    if dur:
+                        note.hold_duration = dur
+                    notes.append(note)
+                    continue
+                n_pts = tokens[i]
+                i += 1
+                positions = []
+                for _ in range(n_pts):
+                    if i >= len(tokens):
+                        break
+                    pt = tokens[i]
+                    if pt in ID_TO_SLD:
+                        positions.append(ID_TO_SLD[pt])
+                    i += 1
+                # Skip SLD_END
+                if i < len(tokens) and tokens[i] == SLD_END_TOKEN:
+                    i += 1
+                # Check for optional duration
+                dur = self._read_dur(tokens, i)
+                if dur:
+                    i += 3  # DUR + beat + subdiv
+                note = TouchNote(beat_div=current_div, positions=positions)
+                note.is_slide = True
+                note.slide_path = list(positions)
+                if dur:
+                    note.hold_duration = dur
+                notes.append(note)
+                continue
+
+            # Simultaneous begin
+            if tid == SIM_BEG:
+                i += 1
+                if i >= len(tokens):
+                    break
+                count_tok = tokens[i]
+                n_notes = 2 if count_tok == SIM_COUNT_2 else int(count_tok)
+                i += 1
+                sub_notes: list[TouchNote] = []
+                dur = None
+                while i < len(tokens) and tokens[i] not in (SIM_END, EOS):
+                    sub_tid = tokens[i]
+                    if sub_tid == DUR:
+                        dur = self._read_dur(tokens, i)
+                        break  # DUR after SIM group
+                    sub_note = self._decode_single_note(sub_tid, current_div)
+                    if sub_note:
+                        sub_notes.append(sub_note)
+                    i += 1
+                if i < len(tokens) and tokens[i] == SIM_END:
+                    i += 1
+                # Merge sub-notes into one simultaneous note
+                if sub_notes:
+                    merged = sub_notes[0]
+                    all_pos = []
+                    has_hold = merged.is_hold
+                    has_break = merged.is_break
+                    is_touch = merged.is_touch
+                    all_touch_regions = list(merged.touch_regions)
+                    for sn in sub_notes:
+                        all_pos.extend(sn.positions)
+                        has_hold = has_hold or sn.is_hold
+                        has_break = has_break or sn.is_break
+                        is_touch = is_touch or sn.is_touch
+                        all_touch_regions.extend(sn.touch_regions)
+                    merged.positions = all_pos
+                    merged.is_simultaneous = True
+                    merged.touch_regions = all_touch_regions
+                    merged.is_touch = is_touch
+                    if dur:
+                        merged.hold_duration = dur
+                        # For touch holds, don't set is_hold
+                        if not is_touch:
+                            has_hold = True
+                    merged.is_hold = has_hold and not is_touch
+                    merged.is_break = has_break and not is_touch
+                    notes.append(merged)
+                continue
+
+            # Duration marker (standalone, should not normally happen)
+            if tid == DUR:
+                i += 3  # skip DUR + 2 values
+                continue
+
+            # Slide end, SIM end (standalone — skip)
+            if tid in (SLD_END_TOKEN, SIM_END):
+                i += 1
+                continue
+
+            # Single note token
+            note = self._decode_single_note(tid, current_div)
+            if note:
+                # Check if next token is DUR (for hold/slide duration)
+                dur = self._read_dur(tokens, i + 1)
+                if dur:
+                    note.hold_duration = dur
+                    # Only set is_hold if not already a slide/touch
+                    if not note.is_slide and not note.is_touch and not note.is_break:
+                        note.is_hold = True
+                    i += 3  # skip DUR + beat + subdiv
+                notes.append(note)
+            i += 1
+
+        from mai_parser.models import Difficulty
+        chart = Chart(difficulty_index=0, difficulty=Difficulty.ReMASTER)
+        chart.notes = notes
+        chart.compute_stats()
+        return chart
+
+    def _decode_single_note(self, tid: int, beat_div: int) -> Optional[TouchNote]:
+        """Decode a single note token (not part of a group)."""
+        note = TouchNote(beat_div=beat_div)
+
+        if tid in ID_TO_TAP:
+            note.positions = [ID_TO_TAP[tid]]
+            return note
+
+        if tid in ID_TO_BRK:
+            note.positions = [ID_TO_BRK[tid]]
+            note.is_break = True
+            return note
+
+        if tid in ID_TO_HLD:
+            note.positions = [ID_TO_HLD[tid]]
+            note.is_hold = True
+            return note
+
+        if tid in ID_TO_SLD:
+            note.positions = [ID_TO_SLD[tid]]
+            note.is_slide = True
+            return note
+
+        if tid in ID_TO_TCH:
+            region = ID_TO_TCH[tid]
+            note.is_touch = True
+            note.touch_regions = [region]
+            return note
+
+        return None
+
+    def _read_dur(self, tokens: list[int], start: int) -> Optional[tuple[int, int]]:
+        """Try to read DUR beat subdiv from tokens[start:]. Returns (beat, subdiv) or None.
+        Clamps to reasonable ranges to filter out hallucinated durations."""
+        return read_duration_tokens(tokens, start)
+
+    # ── Batch ───────────────────────────────────────────────────────
+
+    def encode_batch(self, charts: list[Chart], pad_to: Optional[int] = None,
+                     add_bos: bool = True, add_eos: bool = True,
+                     return_tensors: bool = False):
+        """
+        Encode a batch of charts, padding to the same length.
+
+        Args:
+            charts: List of Chart objects.
+            pad_to: Pad all sequences to this length (auto-detect max if None).
+            add_bos: Prepend BOS.
+            add_eos: Append EOS.
+            return_tensors: If True, return torch.Tensor (requires torch).
+
+        Returns:
+            If return_tensors=False: (list[list[int]], list[int]) = (token_seqs, lengths)
+            If return_tensors=True: (Tensor[batch, max_len], Tensor[batch])
+        """
+        seqs = [self.encode(c, add_bos=add_bos, add_eos=add_eos) for c in charts]
+        lengths = [len(s) for s in seqs]
+        max_len = max(lengths) if pad_to is None else pad_to
+
+        padded = []
+        for seq in seqs:
+            if len(seq) < max_len:
+                seq = seq + [PAD] * (max_len - len(seq))
+            padded.append(seq[:max_len])
+
+        if return_tensors:
+            try:
+                import torch
+                return torch.tensor(padded, dtype=torch.long), torch.tensor(lengths, dtype=torch.long)
+            except ImportError:
+                raise ImportError("torch required for return_tensors=True")
+
+        return padded, lengths
+
+    # ── Debug ───────────────────────────────────────────────────────
+
+    def tokens_to_str(self, tokens: list[int], max_show: int = 60) -> str:
+        """Pretty-print a token sequence with context for raw parameter ids."""
+        parts = []
+        i = 0
+        shown = 0
+        while i < len(tokens) and shown < max_show:
+            tid = tokens[i]
+            if tid == DUR and i + 2 < len(tokens):
+                parts.append("[DUR]")
+                shown += 1
+                if shown < max_show:
+                    parts.append(token_name(tokens[i + 1]))
+                    shown += 1
+                if shown < max_show:
+                    parts.append(token_name(tokens[i + 2]))
+                    shown += 1
+                i += 3
+                continue
+            if tid == SIM_BEG and i + 1 < len(tokens):
+                parts.append("[SIM_BEG]")
+                shown += 1
+                if shown < max_show:
+                    parts.append(token_name(tokens[i + 1]))
+                    shown += 1
+                i += 2
+                continue
+            parts.append(token_name(tid))
+            shown += 1
+            i += 1
+        if i < len(tokens):
+            parts.append(f"... ({len(tokens) - i} more)")
+        return " ".join(parts)
+
+    def print_tokens(self, tokens: list[int], max_show: int = 60) -> None:
+        """Print a token sequence."""
+        print(self.tokens_to_str(tokens, max_show))
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Metadata header builder
+# ═══════════════════════════════════════════════════════════════════════
+
+def build_metadata_header(bpm: float, difficulty: int,
+                          level_value: float, genre: int = 0) -> list[int]:
+    """
+    Build a metadata header token sequence.
+
+    Format: [META_BPM] bpm_byte [META_DIFF] diff [META_LEVEL] level_byte [META_GENRE] genre [META_END]
+
+    This is prepended to chart tokens during training so the model
+    learns to associate metadata with chart style.
+
+    Args:
+        bpm: BPM value (e.g. 173.0)
+        difficulty: 0=BASIC..4=ReMASTER
+        level_value: e.g. 12.4
+        genre: Genre index
+
+    Returns:
+        List of token IDs.
+    """
+    return [
+        META_BPM, int(bpm) // 2,              # BPM 0-510 → 0-255
+        META_DIFF, difficulty,
+        META_LEVEL, min(255, int(level_value * 10)),
+        META_GENRE, genre,
+        META_END,
+    ]
+
+
+def encode_chart_with_header(chart: Chart, bpm: float, difficulty: int,
+                              level_value: float, genre: int = 0) -> list[int]:
+    """Encode chart with grammar-friendly slides (no metadata header, no EOS).
+
+    Metadata (BPM, difficulty, level, genre) is passed as separate condition
+    inputs to the model — NOT as chart tokens. The model learns difficulty
+    from the diff_embed MoE routing, not from token-level metadata.
+
+    HLD_ON/SLD_ON and SLD_MID are inference context/helper tokens, not targets.
+
+    Returns: [BOS] + chart_tokens
+    """
+    tok = MaiChartTokenizer()
+    chart_tokens = tok.encode(chart, add_bos=False, add_eos=False)
+    # add_eos=False avoids appending a synthetic EOS, but parsed charts may
+    # contain a terminal end note. Strip only terminal EOS tokens: raw numeric
+    # values 1/2 are also used as SIM counts, so removing all EOS ids corrupts
+    # simultaneous groups.
+    while chart_tokens and chart_tokens[-1] == EOS:
+        chart_tokens.pop()
+    chart_tokens = unfold_slides(chart_tokens)
+    return [BOS] + chart_tokens
+
+
+def unfold_slides(tokens):
+    """Unfold multi-segment slides into grammar-friendly waypoint tokens.
+
+    SLD_BEG n sld_a sld_b sld_c SLD_END → SLD_BEG sld_a sld_b sld_c SLD_END
+    """
+    result, i = [], 0
+    while i < len(tokens):
+        t = tokens[i]
+        if t == SLD_BEG and i + 2 < len(tokens):
+            n = tokens[i + 1]
+            if 0 < n < 32 and i + 2 + n < len(tokens):
+                pts = tokens[i + 2 : i + 2 + n]
+                result.append(SLD_BEG)
+                result.extend(pts)
+                result.append(SLD_END_TOKEN)
+                i += 2 + n + 1; continue
+        result.append(t); i += 1
+    return result
+
+
+def inject_ongoing_tokens(tokens: list[int]) -> list[int]:
+    """Insert HLD_ON/SLD_ON markers at intermediate positions where a hold/slide is active.
+
+    HLD_n DUR beat subdiv ...tokens... → HLD_ON inserted at each non-DUR position
+    while the hold is active. Same for slides.
+
+    These are informational — the model learns "a hold is ongoing here".
+    During inference they are suppressed; the engine doesn't generate them.
+    """
+    result = []
+    current_div = 4.0
+    hold_beats = 0.0      # remaining beats of active hold
+    slide_beats = 0.0     # remaining beats of active slide
+    dur_skip = 0          # skip DUR parameter tokens
+    i = 0
+
+    while i < len(tokens):
+        t = tokens[i]
+        step = 4.0 / current_div
+
+        # ── Inject ON tokens before note-level tokens ──
+        _is_note = (t >= TAP_BASE and t != DUR) or t == RST or t in ID_TO_DIV
+        if _is_note and dur_skip == 0:
+            if hold_beats > 0:
+                result.append(HLD_ON)
+                hold_beats -= step
+            if slide_beats > 0:
+                result.append(SLD_ON)
+                slide_beats -= step
+
+        # ── Track hold/slide duration ──
+        if dur_skip > 0:
+            dur_skip -= 1
+            result.append(t); i += 1
+            continue
+
+        if t in ID_TO_DIV:
+            current_div = float(ID_TO_DIV.get(t, current_div))
+        elif t == DUR:
+            dur_skip = 2
+        elif t in HLD_TO_ID:
+            # Check if followed by DUR
+            if i + 3 < len(tokens) and tokens[i + 1] == DUR:
+                beat = tokens[i + 2]
+                subdiv = max(tokens[i + 3], 1)
+                hold_beats = beat / subdiv
+        elif t == SLD_BEG:
+            # Find DUR after slide waypoints
+            j = i + 2  # skip SLD_BEG + count
+            while j < len(tokens) and tokens[j] != SLD_END_TOKEN and tokens[j] != DUR:
+                j += 1
+            if j < len(tokens) and tokens[j] == SLD_END_TOKEN:
+                j += 1  # skip SLD_END
+            if j + 2 < len(tokens) and tokens[j] == DUR:
+                beat = tokens[j + 1]
+                subdiv = max(tokens[j + 2], 1)
+                slide_beats = beat / subdiv
+
+        result.append(t)
+        i += 1
+
+    return result
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Chart → maidata text conversion
+# ═══════════════════════════════════════════════════════════════════════
+
+def notes_to_maitext(notes, bpm=150.0):
+    """Convert TouchNote list back to maidata chart text.
+    Format: (173){4}1,2,3h[4:1],5b/8b,
+    """
+    bpm_int = int(bpm)
+    current_div = 4
+    line = f"({bpm_int})"
+    measure = []
+
+    for note in notes:
+        if note.is_end:
+            if measure:
+                line += "{" + str(current_div) + "}" + ",".join(measure) + ","
+            return line + "\nE"
+
+        if note.beat_div != current_div:
+            if measure:
+                line += "{" + str(current_div) + "}" + ",".join(measure) + ","
+                measure = []
+            else:
+                line += "{" + str(current_div) + "},"
+            current_div = note.beat_div
+
+        measure.append(_note_to_text(note))
+
+    if measure:
+        line += "{" + str(current_div) + "}" + ",".join(measure) + ","
+    return line + "\nE"
+
+
+def tokens_to_maitext(tokens, bpm=150.0):
+    """Token sequence → maidata.txt chart text."""
+    tok = MaiChartTokenizer()
+    chart = tok.decode(tokens)
+    return notes_to_maitext(chart.notes, bpm)
+
+
+def _note_to_text(note):
+    """Single TouchNote → maidata text segment."""
+    if note.is_rest:
+        return ""
+    if note.is_touch:
+        # Normalize touch regions: C1..C8 → C, others keep (B7, E2, etc.)
+        regions = []
+        for r in note.touch_regions:
+            if r.startswith("C") and len(r) > 1:
+                regions.append("C")
+            else:
+                regions.append(r)
+        text = "/".join(regions)
+        if note.is_hold and note.hold_duration:
+            text += f"h[{note.hold_duration[0]}:{note.hold_duration[1]}]"
+        return text
+
+    if note.positions:
+        text = "/".join(str(p) for p in note.positions)
+    else:
+        return ""
+
+    if note.is_hold and note.hold_duration:
+        text += f"h[{note.hold_duration[0]}:{note.hold_duration[1]}]"
+    elif note.is_break:
+        text += "b"
+    elif note.is_slide and note.slide_path and len(note.slide_path) >= 2:
+        # Multi-segment slide: all use >
+        start = note.slide_path[0]
+        seg_start = 1
+        while seg_start < len(note.slide_path) and note.slide_path[seg_start] == start:
+            seg_start += 1
+        if seg_start >= len(note.slide_path):
+            return str(start)
+        text = str(start)
+        last_pos = start
+        for p in note.slide_path[seg_start:]:
+            if p != last_pos:
+                text += ">" + str(p)
+                last_pos = p
+        dur = note.hold_duration or (4, 1)
+        text += f"[{dur[0]}:{dur[1]}]"
+    elif note.is_slide and len(note.positions) >= 2:
+        text = str(note.positions[0]) + ">" + str(note.positions[1])
+        dur = note.hold_duration or (4, 1)
+        text += f"[{dur[0]}:{dur[1]}]"
+
+    if note.firework:
+        text += "x"
+    if note.is_star:
+        text += "*"
+    return text