initial BACHI deployment

Dockerfile

@@ -0,0 +1,20 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY packages.txt /app/packages.txt
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git gcc libc6-dev && \
+    xargs apt-get install -y --no-install-recommends < /app/packages.txt && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir gradio==5.28.0
+RUN pip install --no-cache-dir -r /app/requirements.txt
+
+COPY . /app
+
+ENV PORT=7860
+EXPOSE 7860
+
+CMD ["python", "app.py"]

README.md

@@ -0,0 +1,30 @@
+---
+title: Bachi
+emoji: 🎵
+colorFrom: purple
+colorTo: blue
+sdk: docker
+app_port: 7860
+pinned: false
+---
+
+# BACHI — HARP Endpoint
+
+Deploys [BACHI](https://github.com/AndyWeasley2004/BACHI_Chord_Recognition) as a [PyHARP](https://github.com/TEAMuP-dev/pyharp) endpoint.
+
+## What it does
+
+Automatic chord recognition from symbolic music scores. Outputs beat-aligned chord labels (root, quality, bass).
+
+## Inputs
+
+- MIDI or MusicXML file (.mid, .midi, .musicxml, .mxl, .xml)
+- Model type: Classical or Pop
+
+## Note
+
+Model is trained on piano data only and supports MIDI pitch range 21-108.
+
+## License
+
+MIT

app.py

@@ -0,0 +1,236 @@
+import os
+import sys
+import math
+import tempfile
+import pickle
+from pathlib import Path
+
+import torch
+import numpy as np
+import yaml
+import gradio as gr
+from pyharp import ModelCard, build_endpoint
+from huggingface_hub import hf_hub_download
+
+sys.path.insert(0, str(Path(__file__).parent))
+from dataset import load_vocabs
+from models.variants import build_model
+
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+REPO_ID = "Itsuki-music/BACHI_Chord_Recognition"
+CHECKPOINT_NAMES = {
+    "Classical": "classical_film_kdec",
+    "Pop": "pop909_film_kdec",
+}
+CHECKPOINT_FILES = ["best_model.pt", "config.yaml", "vocab.pkl"]
+
+loaded_models = {}
+
+def get_model(model_type: str):
+    if model_type not in loaded_models:
+        folder = CHECKPOINT_NAMES[model_type]
+        ckpt_dir = Path(__file__).parent / "ckpts" / folder
+        ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+        for fname in CHECKPOINT_FILES:
+            dest = ckpt_dir / fname
+            if not dest.exists():
+                print(f"Downloading {folder}/{fname}...", flush=True)
+                hf_hub_download(
+                    repo_id=REPO_ID,
+                    filename=f"{folder}/{fname}",
+                    repo_type="dataset",
+                    local_dir=Path(__file__).parent / "ckpts",
+                )
+                print(f"Downloaded {fname}.", flush=True)
+
+        with open(ckpt_dir / "config.yaml", "r") as f:
+            config = yaml.safe_load(f)
+
+        vocabs = load_vocabs(str(ckpt_dir / "vocab.pkl"))
+        use_key = (
+            bool(config.get("use_key", False))
+            or bool(config["training"].get("use_key", False))
+            or bool(config["model"].get("use_key", False))
+        )
+        experiment = config["experiment"]
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Loading {model_type} model...", flush=True)
+        model = build_model(experiment, config["model"], vocabs, use_key=use_key).to(device)
+        model.load_state_dict(
+            torch.load(ckpt_dir / "best_model.pt", map_location=device, weights_only=True)
+        )
+        model.eval()
+        print(f"{model_type} model loaded.", flush=True)
+
+        loaded_models[model_type] = (model, config, vocabs, use_key, device)
+
+    return loaded_models[model_type]
+
+
+def extract_pianoroll(score_path: Path, resolution: int = 12):
+    import miditoolkit
+    from music21 import converter, note as m21_note, chord as m21_chord
+
+    notes_data = []
+    suffix = score_path.suffix.lower()
+
+    if suffix in {".mid", ".midi"}:
+        midi = miditoolkit.MidiFile(str(score_path))
+        tpb = midi.ticks_per_beat or 480
+        for inst in midi.instruments:
+            if inst.is_drum:
+                continue
+            for n in inst.notes:
+                notes_data.append((n.pitch, n.start / tpb, n.end / tpb))
+    else:
+        sc = converter.parse(str(score_path))
+        parts = list(sc.parts) if sc.hasPartLikeStreams() else [sc]
+        for part in parts:
+            inst = part.getInstrument()
+            if inst:
+                classes = inst.classes if hasattr(inst, "classes") else []
+                if "Percussion" in classes or "Unpitched" in classes:
+                    continue
+            for el in part.flat.notes:
+                dur = float(el.quarterLength)
+                start = float(el.offset)
+                if isinstance(el, m21_note.Note):
+                    notes_data.append((el.pitch.midi, start, start + dur))
+                elif isinstance(el, m21_chord.Chord):
+                    for p in el.pitches:
+                        notes_data.append((p.midi, start, start + dur))
+
+    if not notes_data:
+        return None
+
+    notes_data.sort(key=lambda x: x[1])
+    last_end = max(nd[2] for nd in notes_data)
+    total_frames = math.ceil(last_end * resolution)
+    pianoroll = np.zeros((88, total_frames), dtype=np.int8)
+    for midi_pitch, start_b, end_b in notes_data:
+        row = midi_pitch - 21
+        if not (0 <= row < 88):
+            continue
+        s_f = max(0, math.floor(start_b * resolution))
+        e_f = min(total_frames, math.ceil(end_b * resolution))
+        pianoroll[row, s_f:e_f] = 1
+
+    return pianoroll
+
+
+def predict_piece(pianoroll, model, config, vocabs, use_key, device):
+    beat_resolution = config["model"]["beat_resolution"]
+    label_resolution = config["model"]["label_resolution"]
+    segment_len = config["model"]["n_beats"] * beat_resolution
+    pr_to_label_ratio = beat_resolution // label_resolution
+    comps_eval = ["root", "quality", "bass"] + (["key"] if use_key else [])
+
+    n_frames = pianoroll.shape[0]
+    segments, masks = [], []
+    for i in range(0, n_frames, segment_len):
+        seg = pianoroll[i : i + segment_len]
+        orig_len = seg.shape[0]
+        if orig_len < segment_len:
+            seg = torch.cat([seg, torch.zeros(segment_len - orig_len, seg.shape[1])], dim=0)
+        mask = torch.ones(segment_len, dtype=torch.bool)
+        if orig_len < segment_len:
+            mask[orig_len:] = False
+        segments.append(seg)
+        masks.append(mask)
+
+    piece_preds = {k: [] for k in comps_eval + ["boundary"]}
+    for i in range(0, len(segments), 16):
+        batch_segs = torch.stack(segments[i : i + 16]).to(device)
+        batch_masks = torch.stack(masks[i : i + 16]).to(device)
+        with torch.no_grad():
+            out = model.forward_infer(batch_segs, src_key_padding_mask=~batch_masks)
+        for k in comps_eval + ["boundary"]:
+            if k in out:
+                piece_preds[k].append(out[k].detach().cpu())
+
+    n_target = math.ceil(n_frames / pr_to_label_ratio)
+    piece_pred_ids = {}
+    for k, parts in piece_preds.items():
+        if not parts:
+            continue
+        cat = torch.cat([p.reshape(-1) for p in parts], dim=0)
+        piece_pred_ids[k] = cat[:n_target] if k == "boundary" else cat[:n_target].long()
+
+    inv_root = {v: k for k, v in vocabs["root"].items()}
+    inv_qual = {v: k for k, v in vocabs["quality"].items()}
+    inv_bass = {v: k for k, v in vocabs["bass"].items()}
+
+    valid_len = len(piece_pred_ids.get("root", []))
+    if valid_len == 0:
+        return "No predictions generated."
+
+    r_seq = piece_pred_ids["root"][:valid_len].tolist()
+    q_seq = piece_pred_ids["quality"][:valid_len].tolist()
+    b_seq = piece_pred_ids["bass"][:valid_len].tolist()
+
+    time_per_token = 1.0 / max(1, config["model"]["label_resolution"])
+    merged = []
+    cur_r, cur_q, cur_b = r_seq[0], q_seq[0], b_seq[0]
+    cur_start = 0
+    for t in range(1, valid_len):
+        if r_seq[t] != cur_r or q_seq[t] != cur_q or b_seq[t] != cur_b:
+            label = f"{inv_root.get(cur_r, str(cur_r))}_{inv_qual.get(cur_q, str(cur_q))}_{inv_bass.get(cur_b, str(cur_b))}"
+            merged.append(f"{cur_start * time_per_token:.2f} {label}")
+            cur_r, cur_q, cur_b = r_seq[t], q_seq[t], b_seq[t]
+            cur_start = t
+    label = f"{inv_root.get(cur_r, str(cur_r))}_{inv_qual.get(cur_q, str(cur_q))}_{inv_bass.get(cur_b, str(cur_b))}"
+    merged.append(f"{cur_start * time_per_token:.2f} {label}")
+
+    return "\n".join(merged)
+
+
+model_card = ModelCard(
+    name="BACHI Chord Recognition",
+    description="Automatic chord recognition from symbolic music scores (MIDI or MusicXML). Outputs beat-aligned chord labels.",
+    author="Mingyang Yao, Ke Chen, Shlomo Dubnov, Taylor Berg-Kirkpatrick",
+    tags=["chord-recognition", "symbolic-music", "midi", "musicxml"],
+)
+
+
+def process_fn(input_file: str, model_type: str) -> str:
+    print(f"Processing with {model_type} model...", flush=True)
+    model, config, vocabs, use_key, device = get_model(model_type)
+
+    score_path = Path(input_file)
+    pianoroll_np = extract_pianoroll(score_path, resolution=config["model"]["beat_resolution"])
+    if pianoroll_np is None:
+        return "Error: Could not extract notes from the input file."
+
+    pianoroll = torch.from_numpy(pianoroll_np.T).float()
+    result = predict_piece(pianoroll, model, config, vocabs, use_key, device)
+    print("Done.", flush=True)
+    return result
+
+
+with gr.Blocks() as demo:
+    input_components = [
+        gr.File(
+            label="Input Score (.mid, .midi, .musicxml, .mxl, .xml)",
+            file_types=[".mid", ".midi", ".musicxml", ".mxl", ".xml"],
+        ),
+        gr.Dropdown(
+            choices=["Classical", "Pop"],
+            value="Classical",
+            label="Model Type",
+        ),
+    ]
+    output_components = [
+        gr.Textbox(label="Chord Predictions", lines=20),
+    ]
+    app = build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+print("Launching Gradio...", flush=True)
+demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)

dataset.py

@@ -0,0 +1,382 @@
+import os
+import json
+import torch
+from torch.utils.data import Dataset, random_split
+import numpy as np
+import pickle
+from typing import List, Dict, Tuple, Optional, Any
+import math
+import torch.nn.functional as F
+import random
+from collections import defaultdict
+
+
+def load_vocabs(vocab_path: str) -> Dict[str, Any]:
+    """Loads vocabularies and augments with per-component PAD/NONE indices.
+
+    For 3-part prediction, only `root`, `quality`, and `bass` are loaded.
+    """
+    with open(vocab_path, 'rb') as f:
+        data = pickle.load(f)
+
+    root_map = data['root_to_idx']
+    pad_token = 'PAD'
+    none_tokens = ['N', 'None']  # allow either spelling in source vocabs
+    bass_map = root_map
+
+    # Only keep the three chord parts for prediction
+    vocabs = {
+        'root': root_map,
+        'quality': data['quality_to_idx'],
+        'bass': bass_map,
+        'key': data['key_to_idx'],
+    }
+
+    # Global root PAD index (back-compat)
+    vocabs['pad_idx'] = root_map[pad_token]
+
+    # Add per-component PAD and NONE indices
+    for comp, comp_map in list(vocabs.items()):
+        if comp == 'pad_idx':
+            continue
+        # per-component PAD index (must exist)
+        comp_pad_idx = comp_map.get(pad_token)
+        if comp_pad_idx is None:
+            raise ValueError(f"Component '{comp}' vocab lacks PAD token")
+        vocabs[f'{comp}_pad_idx'] = comp_pad_idx
+
+        # NONE index preference: N > None > PAD
+        none_idx = None
+        for tok in none_tokens:
+            if tok in comp_map:
+                none_idx = comp_map[tok]
+                break
+        if none_idx is None:
+            none_idx = comp_pad_idx
+        vocabs[f'{comp}_none_idx'] = none_idx
+
+    return vocabs
+
+
+class PianoRollDataset(Dataset):
+    """Dataset for piano roll representation."""
+    pad_idx = -1 # Will be updated in __init__
+    def __init__(
+        self,
+        data_root: str,
+        config: dict,
+        vocabs: Dict[str, Any],
+        split: str = 'train',
+        use_augmentation: bool = False,
+        use_key: bool = False,
+    ):
+        self.data_root = data_root
+        self.config = config
+        self.n_beats = self.config['n_beats']
+        self.split = split
+        self.use_augmentation = use_augmentation
+        self.use_key = use_key
+        self.beat_resolution = self.config['beat_resolution']
+        self.label_resolution = self.config['label_resolution']
+        self.pr_to_label_ratio = self.beat_resolution // self.label_resolution
+
+        self.vocabs = vocabs
+        self.pad_idx = self.vocabs['pad_idx']
+
+        self.chord_components = ['root', 'quality', 'bass']
+        self.label_indices_map = {'root': 0, 'quality': 1, 'bass': 2}
+        if self.use_key:
+            self.chord_components.append('key')
+            self.label_indices_map['key'] = 3
+
+        # --- Lengths in pianoroll-frame resolution ---
+        self.max_len = self.n_beats * self.beat_resolution
+
+        for comp in self.chord_components:
+            setattr(self, f'{comp}_vocab', self.vocabs[comp])
+            setattr(self, f'{comp}_none_idx', self.vocabs[f'{comp}_none_idx'])
+
+        suffix = 'shift0.npz' if not self.use_augmentation else '.npz'
+        # print(f"Loading {suffix} files from {data_root}")
+        self.file_list = sorted([
+            os.path.join(data_root, f) 
+            for f in os.listdir(data_root) if f.endswith(suffix)
+        ])
+
+    def __len__(self) -> int:
+        return len(self.file_list)
+
+    def __getitem__(self, idx: int) -> Optional[Dict[str, torch.Tensor]]:
+        filepath = self.file_list[idx]
+        with np.load(filepath, allow_pickle=True) as data:
+            pianoroll_full = torch.from_numpy(data['pianoroll'].T).float()
+            labels_full = data['labels']
+            boundaries_full = data['boundaries']
+ 
+        pianoroll = pianoroll_full
+        labels = labels_full
+
+        # --- Create ground truth chord tensor from labels (map to per-component vocab indices) ---
+        target_indices = {}
+
+        for comp in self.chord_components:
+            vocab = getattr(self, f'{comp}_vocab')
+            none_idx = getattr(self, f'{comp}_none_idx')
+            label_col_idx = self.label_indices_map[comp]
+            col = labels[:, label_col_idx]
+            mapped_tensor = None
+            # If labels are already integer indices within range, accept directly
+            try:
+                if np.issubdtype(col.dtype, np.integer):
+                    col_int = col.astype(np.int64)
+                    if col_int.min(initial=0) >= 0 and col_int.max(initial=0) < len(vocab):
+                        mapped_tensor = torch.from_numpy(col_int)
+            except Exception:
+                mapped_tensor = None
+            # Otherwise map string/mixed labels through vocab with fallback to none_idx
+            if mapped_tensor is None:
+                try:
+                    col_list = col.astype(str).tolist()
+                except Exception:
+                    col_list = [str(x) for x in col.tolist()]
+                mapped = [vocab.get(lbl, none_idx) for lbl in col_list]
+                mapped_tensor = torch.tensor(mapped, dtype=torch.long)
+            target_indices[comp] = mapped_tensor.long()
+
+        # --- Load pre-computed boundary flag ---
+        boundary_flag = torch.from_numpy(boundaries_full.astype(np.float32))
+        
+        if self.split == 'train':
+            return self._get_train_item(pianoroll, target_indices, boundary_flag)
+        else: # 'val' or 'test'
+            piece_name = _get_piece_name(filepath)
+            # Build accurate targets from labels for evaluation
+            return self._get_eval_item(pianoroll, labels, boundary_flag, piece_name)
+    
+    def _sample_stratified_start(self, X: int) -> int:
+        """
+        Sample s ∈ {0..X} with P(s) ∝ 1 + beta * (s/X).
+        Implemented as a mixture of Uniform and 'linear-in-s' discrete law.
+        Exact, O(1), numerically stable.
+
+        beta ∈ [0,2]. beta=0 -> uniform; beta=1 -> mild late tilt (good default).
+        """
+        if X <= 0:
+            return 0
+
+        beta = float(getattr(self, 'stratify_beta', 1.0))
+
+        # Mixture weights: P = a * Uniform + (1-a) * Linear(s)
+        a = 1.0 - beta / 2.0  # ∈ [0,1]
+        if np.random.rand() < a:
+            # Uniform over 0..X
+            return int(np.random.randint(0, X + 1))
+        else:
+            # Sample from Q(s) ∝ s over {0..X} (i.e., s=0 has weight 0).
+            # Do it by inverting triangular numbers over 1..X.
+            M = X * (X + 1) // 2  # sum_{s=1}^X s
+            r = np.random.randint(1, M + 1)  # 1..M inclusive
+            s = int((math.isqrt(1 + 8 * r) - 1) // 2)  # floor((sqrt(1+8r)-1)/2)
+            # Numerical guard (rare when r hits exact triangle): clamp
+            if s > X:
+                s = X
+            return s
+
+    def _get_train_item(self, pianoroll, target_indices, boundary_flag):
+        n_pr_frames = pianoroll.shape[0]
+        # start with at least half of window size and convert to label frames
+        max_start_label_frame = (n_pr_frames - self.max_len // 2) // self.pr_to_label_ratio
+        target_max_len = self.max_len // self.pr_to_label_ratio
+        
+        # Stratified start over 0..max_start_label_frame (tilt to late positions)
+        start_label_frame = self._sample_stratified_start(max_start_label_frame)
+        start_pr_frame = start_label_frame * self.pr_to_label_ratio
+        
+        # --- slice & pad encoder input ---
+        pr_segment = pianoroll[start_pr_frame : start_pr_frame + self.max_len]
+        pr_pad_amount = self.max_len - pr_segment.shape[0]
+        if pr_pad_amount > 0:
+            # keep dtype/device consistent with pr_segment
+            pr_pad = pr_segment.new_zeros((pr_pad_amount, pr_segment.shape[1]))
+            pr_segment = torch.cat([pr_segment, pr_pad], dim=0)
+
+        # --- slice targets at label resolution ---
+        target_start = start_label_frame
+        target_segs = {}
+        for comp in self.chord_components:
+            target_segs[comp] = target_indices[comp][target_start : target_start + target_max_len]
+        boundary_seg = boundary_flag[target_start : target_start + target_max_len]
+
+        # --- masks & padding for targets ---
+        current_target_len = target_segs[self.chord_components[0]].shape[0]
+        target_mask = torch.zeros(target_max_len, dtype=torch.bool)
+        target_mask[:current_target_len] = True
+
+        # expand target mask to encoder (frame) mask
+        encoder_mask = target_mask.repeat_interleave(self.pr_to_label_ratio)
+        if pr_pad_amount > 0:
+            encoder_mask[-pr_pad_amount:] = False
+
+        target_pad_amount = target_max_len - current_target_len
+        if target_pad_amount > 0:
+            for comp in self.chord_components:
+                comp_none_idx = getattr(self, f'{comp}_none_idx')
+                pad_tensor = torch.full((target_pad_amount,), comp_none_idx, dtype=torch.long)
+                target_segs[comp] = torch.cat([target_segs[comp], pad_tensor])
+
+            boundary_pad = torch.zeros(target_pad_amount, dtype=boundary_seg.dtype)
+            boundary_seg = torch.cat([boundary_seg, boundary_pad])
+
+        item = {
+            'encoder_input': pr_segment,
+            'target_boundary': boundary_seg,
+            'mask': target_mask,
+            'encoder_mask': encoder_mask,
+        }
+        for comp in self.chord_components:
+            item[f'target_{comp}'] = target_segs[comp]
+
+        return item
+
+    def _get_eval_item(self, pianoroll, labels, boundary_flag, piece_name):
+        # Reconstruct per-component target indices directly from the label matrix
+        n_label_frames = labels.shape[0]
+        target_indices = {}
+        for comp in self.chord_components:
+            vocab = getattr(self, f'{comp}_vocab')
+            none_idx = getattr(self, f'{comp}_none_idx')
+            label_col_idx = self.label_indices_map[comp]
+            # Extract the column for this component; handle types robustly
+            col = labels[:, label_col_idx]
+            mapped_tensor = None
+            # Case 1: already integer indices
+            try:
+                if np.issubdtype(col.dtype, np.integer):
+                    col_int = col.astype(np.int64)
+                    # If values look like valid indices, accept directly; otherwise fallback to mapping
+                    if col_int.min(initial=0) >= 0 and col_int.max(initial=0) < len(vocab):
+                        mapped_tensor = torch.from_numpy(col_int)
+            except Exception:
+                mapped_tensor = None
+            # Case 2: map from labels (strings or mixed types) to indices
+            if mapped_tensor is None:
+                try:
+                    col_list = col.astype(str).tolist()
+                except Exception:
+                    col_list = [str(x) for x in col.tolist()]
+                mapped = [vocab.get(lbl, none_idx) for lbl in col_list]
+                mapped_tensor = torch.tensor(mapped, dtype=torch.long)
+            target_indices[comp] = mapped_tensor.long()
+
+        mask = torch.ones(n_label_frames, dtype=torch.bool)
+        encoder_mask = torch.ones(pianoroll.shape[0], dtype=torch.bool)
+        item = {
+            'piece_name': piece_name,
+            'encoder_input': pianoroll,
+            'target_boundary': boundary_flag,
+            'mask': mask,
+            'encoder_mask': encoder_mask,
+        }
+        for comp in self.chord_components:
+            item[f'target_{comp}'] = target_indices[comp]
+        return item
+
+    def get_vocab_sizes(self) -> Dict[str, int]:
+        sizes = {comp: len(self.vocabs[comp]) for comp in self.chord_components}
+        return sizes
+
+    def get_pad_idx(self) -> int:
+        return self.pad_idx
+
+
+def _get_piece_name(filename: str) -> str:
+    """Extracts the base piece name from a filename by splitting on '_shift'."""
+    base_filename = os.path.basename(filename)
+    if '_shift' in base_filename:
+        piece_name = base_filename.split('_shift')[0]
+    else:
+        piece_name = base_filename
+    return piece_name
+
+
+def create_datasets(
+    data_root: str,
+    config: dict,
+    vocabs: Dict[str, Any],
+    seed: int = 42,
+) -> Tuple[Dataset, Dataset]:
+    """
+    Create train and validation datasets with group-based splitting.
+    This ensures that all augmentations of a piece belong to the same split.
+    """
+    full_dataset = PianoRollDataset(
+        data_root=data_root,
+        config=config,
+        vocabs=vocabs,
+        split='train',  # split does not matter here
+        use_augmentation=config['use_augmentation'],
+        use_key=config['use_key'],
+    )
+
+    # Group files by piece name
+    piece_files = defaultdict(list)
+    for f in full_dataset.file_list:
+        piece_name = _get_piece_name(f)
+        piece_files[piece_name].append(f)
+
+    unique_pieces = sorted(list(piece_files.keys()))
+
+    # Shuffle for random split
+    random.seed(seed)
+    random.shuffle(unique_pieces)
+
+    # Split unique pieces (90% train, 10% validation)
+    train_size = int(0.9 * len(unique_pieces))
+    train_pieces = unique_pieces[:train_size]
+    val_pieces = unique_pieces[train_size:]
+
+    # Get file lists for each split, only use shift0.npz for validation
+    train_files = [file for piece in train_pieces for file in piece_files[piece]]
+    if config['use_augmentation']:
+        val_files = [file for piece in val_pieces for file in piece_files[piece] if file.endswith('shift0.npz')]
+    else:
+        val_files = [file for piece in val_pieces for file in piece_files[piece]]
+    print(f"Train files: {len(train_files)}, Val files: {len(val_files)}")
+
+    # Create datasets for each split with the correct file list
+    train_dataset = PianoRollDataset(data_root, config, vocabs, 'train', use_key=config['use_key'])
+    train_dataset.file_list = train_files
+
+    val_dataset = PianoRollDataset(data_root, config, vocabs, 'val', use_key=config['use_key'])
+    val_dataset.file_list = val_files
+
+    json.dump(sorted([_get_piece_name(file) for file in val_files]), 
+              open('val_files_unique.json', 'w'), indent=2)
+
+    return train_dataset, val_dataset
+
+
+def collate_fn(batch):
+    """
+    Collate function that filters out empty or invalid samples.
+    For training, it uses default collate.
+    For evaluation (variable length), it handles padding if needed, but typically used with batch_size=1.
+    """
+    batch = [item for item in batch if item is not None]
+    if not batch:
+        return {}
+
+    # If batch contains only a single sample, simply return that sample's dict.
+    # This is handy for evaluation where we usually set batch_size = 1 and do
+    # not need the extra list wrapper.
+    if len(batch) == 1 and 'piece_name' in batch[0]:
+        return batch[0]
+
+    # For training batches (fixed-length segments) every sample has the same
+    # sequence length, so the default PyTorch collate works fine.
+    if 'encoder_input' in batch[0] and batch[0]['encoder_input'].shape[0] == batch[-1]['encoder_input'].shape[0]:
+        return torch.utils.data.dataloader.default_collate(batch)
+
+    # Otherwise we have variable-length sequences – fall back to returning the
+    # list so the caller can deal with padding/iteration manually.
+    return batch

models/HT.py

@@ -0,0 +1,386 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Dict, Optional, Tuple
+
+def get_relative_position_encoding(n_steps, d_model, max_dist=10):
+    """
+    Generates relative positional encodings, similar to Transformer-XL.
+    """
+    vocab_size = 2 * max_dist + 1
+    position = torch.arange(0, vocab_size, dtype=torch.float).unsqueeze(1)
+    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+    
+    pe = torch.zeros(vocab_size, d_model)
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    
+    range_vec = torch.arange(n_steps)
+    distance_mat = range_vec[None, :] - range_vec[:, None]
+    distance_mat_clipped = torch.clamp(distance_mat, -max_dist, max_dist)
+    final_mat = distance_mat_clipped + max_dist
+
+    embeddings = F.embedding(final_mat.long(), pe)
+    return embeddings
+
+class RelativeMultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention with relative positional encoding, inspired by Transformer-XL and HTv2.
+    """
+    def __init__(self, d_model, n_heads, dropout=0.1, max_dist=10):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads
+        self.max_dist = max_dist
+
+        self.w_q = nn.Linear(d_model, d_model)
+        self.w_k = nn.Linear(d_model, d_model)
+        self.w_v = nn.Linear(d_model, d_model)
+        self.w_o = nn.Linear(d_model, d_model)
+        
+        self.u_bias = nn.Parameter(torch.randn(self.n_heads, self.d_head))
+        self.v_bias = nn.Parameter(torch.randn(self.n_heads, self.d_head))
+        self.w_r = nn.Linear(d_model, d_model)
+        
+        self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(d_model)
+
+    def forward(self, q, k, v, pos_emb, key_padding_mask=None, attn_mask=None, values_provided=False):
+        batch_size, seq_len_q, _ = q.size()
+        seq_len_k = k.size(1)
+
+        residual = q
+
+        q = self.w_q(q).view(batch_size, seq_len_q, self.n_heads, self.d_head)
+        k = self.w_k(k).view(batch_size, seq_len_k, self.n_heads, self.d_head)
+        v_transformed = self.w_v(v).view(batch_size, seq_len_k, self.n_heads, self.d_head)
+        
+        q_with_u = (q + self.u_bias).transpose(1, 2)  # (B, h, T_q, d_h)
+        q_with_v = (q + self.v_bias).transpose(1, 2)  # (B, h, T_q, d_h)
+        k = k.transpose(1, 2) # (B, h, T_k, d_h)
+        v_transformed = v_transformed.transpose(1, 2) # (B, h, T_k, d_h)
+
+        pos_emb = self.w_r(pos_emb).view(seq_len_q, seq_len_k, self.n_heads, self.d_head)
+        pos_emb = pos_emb.permute(2, 0, 3, 1) # (h, T_q, d_h, T_k)
+        
+        # Content-based addressing
+        ac = torch.matmul(q_with_u, k.transpose(-2, -1)) # (B, h, T_q, T_k)
+        
+        # Position-based addressing
+        q_for_bd = q_with_v.permute(1, 2, 0, 3) # (h, T_q, B, d_h)
+        bd_t = torch.matmul(q_for_bd, pos_emb) # (h, T_q, B, T_k)
+        bd = bd_t.permute(2, 0, 1, 3) # (B, h, T_q, T_k)
+
+        attn_score = (ac + bd) / math.sqrt(self.d_head)
+
+        # Apply key and (optional) attention masks
+        if key_padding_mask is not None:
+            attn_score = attn_score.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+
+        if attn_mask is not None:
+            attn_score = attn_score.masked_fill(attn_mask.unsqueeze(0), float('-inf'))
+
+        # ------------------------------------------------------------------
+        # Guard against rows where *all* keys were masked (e.g., an attention
+        # block that consists purely of padding tokens). A softmax over a
+        # vector of \(-\infty\) would otherwise produce NaNs. We detect such
+        # rows and set their scores to zero, ensuring a well-defined softmax
+        # (which will yield a uniform distribution) before later resetting
+        # their weights to zero.
+        # ------------------------------------------------------------------
+        all_masked = torch.isinf(attn_score) & (attn_score < 0)  # True where -inf
+        all_masked = all_masked.all(dim=-1, keepdim=True)        # (B, h, T_q, 1)
+
+        # For rows with all_masked == True, replace -inf with 0 so that the
+        # subsequent softmax does not generate NaNs.
+        attn_score = attn_score.masked_fill(all_masked, 0.0)
+        
+        attn_weights = F.softmax(attn_score, dim=-1)  # (B, h, T_q, T_k)
+
+        # After the softmax, zero-out rows that correspond to fully-padded
+        # queries so they don’t influence the output.
+        attn_weights = attn_weights.masked_fill(all_masked, 0.0)
+        attn_weights = self.dropout(attn_weights)
+
+        attn_output = torch.matmul(attn_weights, v_transformed) # (B, h, T_q, d_h)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, self.d_model)
+        
+        attn_output = self.w_o(attn_output)
+        attn_output = self.dropout(attn_output)
+
+        if values_provided:
+            output = v + attn_output
+        else:
+            output = residual + attn_output
+        
+        return self.layer_norm(output), attn_weights
+
+class IntraBlockMHA(nn.Module):
+    def __init__(self, d_model, n_heads, dropout, max_dist=3):
+        super().__init__()
+        self.mha = RelativeMultiHeadAttention(d_model, n_heads, dropout, max_dist=max_dist)
+    
+    def forward(self, x, n_blocks, key_padding_mask=None):
+        batch_size, seq_len, _ = x.shape
+        block_len = seq_len // n_blocks
+        
+        x = x.reshape(batch_size * n_blocks, block_len, -1)
+        
+        mask = None
+        if key_padding_mask is not None:
+            mask = key_padding_mask.reshape(batch_size * n_blocks, block_len)
+
+        pos_emb = get_relative_position_encoding(block_len, x.size(-1), self.mha.max_dist).to(x.device)
+
+        x, _ = self.mha(x, x, x, pos_emb, key_padding_mask=mask)
+        return x.reshape(batch_size, seq_len, -1)
+
+class ConvFFN(nn.Module):
+    def __init__(self, d_model, dropout=0.1):
+        super().__init__()
+        # piano roll resolution: 4 -> 12 (3x)
+        # kernel size: 3 -> 9 (3x)
+        self.conv1 = nn.Conv1d(d_model, d_model, kernel_size=9, padding=4)
+        self.conv2 = nn.Conv1d(d_model, d_model, kernel_size=9, padding=4)
+        self.dropout = nn.Dropout(dropout)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        residual = x
+        x = x.transpose(1, 2)
+        x = F.relu(self.conv1(x))
+        x = self.dropout(x)
+        x = self.conv2(x)
+        x = self.dropout(x)
+        x = x.transpose(1, 2)
+        x = x + residual
+        return self.norm(x)
+
+class TransformerLayer(nn.Module):
+    def __init__(self, d_model, n_heads, dropout, max_dist, use_conv_ffn=True):
+        super().__init__()
+        self.self_attn = RelativeMultiHeadAttention(d_model, n_heads, dropout, max_dist)
+        self.cross_attn = RelativeMultiHeadAttention(d_model, n_heads, dropout, max_dist)
+        self.pos_attn = RelativeMultiHeadAttention(d_model, n_heads, dropout, max_dist)
+        if use_conv_ffn:
+            self.ffn = ConvFFN(d_model, dropout)
+        else: # Fallback to original FFN if needed
+            self.ffn = nn.Sequential(
+                nn.Linear(d_model, d_model * 4),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(d_model * 4, d_model),
+                nn.Dropout(dropout)
+            )
+            self.norm = nn.LayerNorm(d_model)
+
+        self.use_conv_ffn = use_conv_ffn
+
+    def forward(self, dec_input, enc_output, pos_emb, dec_pos_emb,
+                tgt_mask=None, memory_mask=None,
+                tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        # Self-attention
+        dec_output, _ = self.self_attn(dec_input, dec_input, dec_input, pos_emb, key_padding_mask=tgt_key_padding_mask)
+
+        # Positional Attention
+        dec_output = self.pos_attn(dec_pos_emb.unsqueeze(0).repeat(dec_input.size(0), 1, 1), 
+                                   dec_pos_emb.unsqueeze(0).repeat(dec_input.size(0), 1, 1), 
+                                   dec_output, pos_emb=pos_emb, 
+                                   key_padding_mask=tgt_key_padding_mask, values_provided=True)[0]
+
+        # Cross-attention
+        dec_output, _ = self.cross_attn(dec_output, enc_output, enc_output, pos_emb, key_padding_mask=memory_key_padding_mask)
+
+        # FFN
+        if self.use_conv_ffn:
+            dec_output = self.ffn(dec_output)
+        else:
+            dec_output = self.norm(dec_output + self.ffn(dec_output))
+            
+        return dec_output
+
+class PositionalEncoding(nn.Module):
+    """Injects some information about the relative or absolute position of the tokens in the sequence."""
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        """
+        Args:
+            x: Tensor, shape [batch_size, seq_len, d_model]
+        """
+        x = x + self.pe[:, :x.size(1), :]
+        return self.dropout(x)
+
+class BinaryRound(torch.autograd.Function):
+    """
+    Rounds a tensor whose values are in [0,1] to a tensor with values in {0, 1},
+    using the straight through estimator for the gradient.
+    """
+    @staticmethod
+    def forward(ctx, input):
+        return torch.round(input).to(input.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+class HarmonyTransformer(nn.Module):
+    def __init__(self, config: Dict, vocab_sizes: Dict[str, int]):
+        super().__init__()
+        self.input_size = config['input_size']
+        self.d_model = config['d_model']
+        self.n_layers = config['n_layers']
+        self.n_heads = config['n_heads']
+        self.dropout_rate = config['dropout']
+        self.train_boundary = config['train_boundary']
+        self.slope = config.get('slope', 1.0)
+        self.max_len = config['n_beats'] * config['beat_resolution']
+        self.config = config
+        
+        # Embeddings
+        self.enc_input_embed = nn.Linear(self.input_size, self.d_model)
+        self.dec_input_embed = nn.Linear(self.input_size, self.d_model)
+        
+        self.pos_encoder = PositionalEncoding(self.d_model, self.dropout_rate, self.max_len)
+        self.register_buffer('pos_emb', get_relative_position_encoding(self.max_len, self.d_model, self.max_len - 1))
+        
+        self.enc_intra_block_mha = IntraBlockMHA(self.d_model, self.n_heads, self.dropout_rate, max_dist=3)
+        self.dec_intra_block_mha = IntraBlockMHA(self.d_model, self.n_heads, self.dropout_rate, max_dist=3)
+
+        # Encoder
+        self.encoder_layers = nn.ModuleList([
+            RelativeMultiHeadAttention(self.d_model, self.n_heads, self.dropout_rate, self.max_len-1)
+            for _ in range(self.n_layers)
+        ])
+        self.encoder_ffns = nn.ModuleList([ConvFFN(self.d_model, self.dropout_rate) for _ in range(self.n_layers)])
+        self.enc_weights = nn.Parameter(torch.zeros(self.n_layers + 1))
+
+        # Decoder
+        self.decoder_layers = nn.ModuleList([
+            TransformerLayer(self.d_model, self.n_heads, self.dropout_rate, self.max_len-1)
+            for _ in range(self.n_layers)
+        ])
+        self.dec_weights = nn.Parameter(torch.zeros(self.n_layers + 1))
+
+        # Chord Change Prediction
+        self.chord_change_predictor = nn.Linear(self.d_model, 1)
+
+        # Output layers
+        self.root_predictor = nn.Linear(self.d_model, vocab_sizes['root'])
+        self.quality_predictor = nn.Linear(self.d_model, vocab_sizes['quality'])
+        self.bass_predictor = nn.Linear(self.d_model, vocab_sizes['bass'])
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def chord_block_compression(self, hidden_states, chord_changes):
+        """Compress hidden states according to chord changes."""
+        block_ids = torch.cumsum(chord_changes, dim=1) - chord_changes[:, 0].unsqueeze(1)
+        # Ensure integer dtype for one-hot encoding
+        block_ids = block_ids.long()
+
+        max_blocks = (torch.max(block_ids).item() + 1) if block_ids.numel() > 0 else 1
+        one_hot_ids = F.one_hot(block_ids, num_classes=max_blocks).float()  # (B, S, M)
+        
+        summed_states = torch.bmm(one_hot_ids.transpose(1, 2), hidden_states)  # (B, M, H)
+        block_counts = one_hot_ids.sum(dim=1).unsqueeze(-1).clamp(min=1)
+        
+        mean_states = summed_states / block_counts
+        # block_ids already of integer dtype
+        return mean_states, block_ids
+
+    def decode_compressed_sequences(self, compressed_sequences, block_ids):
+        """Decode chord sequences according to chords_pred and block_ids."""
+        return torch.gather(compressed_sequences, 1, block_ids.unsqueeze(-1).expand(-1, -1, compressed_sequences.size(-1)))
+
+    def forward(self, src: torch.Tensor, src_key_padding_mask: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            src: (batch_size, seq_len, input_size)
+            src_key_padding_mask: (batch_size, seq_len), True for valid, False for pad
+        Returns:
+            Dictionary of predictions.
+        """
+        # --- Encoder ---
+        enc_output = self.pos_encoder(self.enc_input_embed(src))
+        
+        # Intra-block MHA
+        enc_output = self.enc_intra_block_mha(enc_output, n_blocks=self.max_len//4, key_padding_mask=src_key_padding_mask)
+
+        # Main encoder layers
+        enc_layer_outputs = [enc_output]
+        for i in range(self.n_layers):
+            enc_output, _ = self.encoder_layers[i](enc_output, enc_output, enc_output, 
+                                                   self.pos_emb, key_padding_mask=src_key_padding_mask)
+            enc_output = self.encoder_ffns[i](enc_output)
+            enc_layer_outputs.append(enc_output)
+
+        # Weighted sum of encoder layers
+        enc_weights = F.softmax(self.enc_weights, dim=0)
+        enc_output = torch.stack(enc_layer_outputs, dim=-1)  # (B, S, H, L+1)
+        enc_output = (enc_output * enc_weights).sum(dim=-1)  # (B, S, H)
+
+        # --- Chord‐change prediction ---
+        boundary_logits = self.chord_change_predictor(enc_output)  # (B, S, 1)
+        chord_change_prob = torch.sigmoid(self.slope * boundary_logits)
+        chord_change_pred = BinaryRound.apply(chord_change_prob).squeeze(-1)  # (B, S) in {0,1}
+
+        # --- Decoder input embedding with regionalization ---
+        dec_input_embed = self.dec_input_embed(src)
+        dec_input_embed = F.dropout(dec_input_embed, p=self.dropout_rate, training=self.training)
+        dec_input_embed = self.dec_intra_block_mha(dec_input_embed, n_blocks=self.max_len // 4, key_padding_mask=src_key_padding_mask)
+
+        # Compress by predicted chord boundaries and expand back
+        dec_input_embed_reg, block_ids = self.chord_block_compression(dec_input_embed, chord_change_pred.long())
+        dec_input_embed_reg = self.decode_compressed_sequences(dec_input_embed_reg, block_ids)
+
+        # Combine embeddings
+        dec_input_embed = dec_input_embed + dec_input_embed_reg + enc_output
+
+        # Positional encoding
+        dec_input_embed = self.pos_encoder(dec_input_embed)
+        dec_pos_emb = self.pos_encoder.pe[:, :dec_input_embed.size(1), :].squeeze(0)
+
+        # --- Decoder layers with layer weighting ---
+        dec_layer_outputs = [dec_input_embed]
+        dec_output = dec_input_embed
+        for i in range(self.n_layers):
+            dec_output = self.decoder_layers[i](dec_output, enc_output, self.pos_emb, dec_pos_emb,
+                                                tgt_key_padding_mask=src_key_padding_mask,
+                                                memory_key_padding_mask=src_key_padding_mask)
+            dec_layer_outputs.append(dec_output)
+
+        dec_weights = F.softmax(self.dec_weights, dim=0)
+        dec_output = torch.stack(dec_layer_outputs, dim=-1)  # (B, S, H, L+1)
+        dec_output = (dec_output * dec_weights).sum(dim=-1)
+
+        # --- Output Projections ---
+        root_logits = self.root_predictor(dec_output)
+        quality_logits = self.quality_predictor(dec_output)
+        bass_logits = self.bass_predictor(dec_output)
+
+        preds = {
+            'root': root_logits,
+            'quality': quality_logits,
+            'bass': bass_logits,
+            'boundary': boundary_logits
+        }
+
+        return preds

models/components.py

@@ -0,0 +1,339 @@
+import math
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, d_model: int, frames_per_patch: int = 6, expansion: int = 2):
+        super().__init__()
+        self.d_model = d_model
+        self.frames_per_patch = frames_per_patch
+        # Frame embedding (collapse pitch dim)
+        self.conv2d = nn.Conv2d(
+            in_channels=1,
+            out_channels=d_model,
+            kernel_size=(88, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+        )
+        self.norm_frame = nn.LayerNorm(d_model)
+        # anti-aliasing conv on time axis
+        self.aa = nn.Conv1d(d_model, d_model, kernel_size=3, stride=1,
+                            padding=1, groups=d_model, bias=False)
+        
+        # Late temporal pooling (downsample frames -> patches)
+        self.glu_conv = nn.Conv1d(
+            in_channels=d_model,
+            out_channels=d_model * expansion * 2,
+            kernel_size=frames_per_patch,
+            stride=frames_per_patch,
+            padding=0,
+            bias=True,
+        )
+        self.project = nn.Conv1d(
+            in_channels=d_model * expansion,
+            out_channels=d_model,
+            kernel_size=1,
+        )
+        self.norm_temporal = nn.LayerNorm(d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, 1, 88, T)
+        x = self.conv2d(x)                      # (B, C, 1, T)
+        x = x.squeeze(2).transpose(1, 2)        # (B, T, C)
+        x = self.norm_frame(x)
+
+        # anti-aliased and temporal pooling
+        x = x.transpose(1, 2)                   # (B, C, T)
+        x = self.aa(x)                          # (B, C, T)
+        v, g = self.glu_conv(x).chunk(2, dim=1)
+        x = self.project(v * torch.sigmoid(g))  # (B, C, T//k)
+        x = x.transpose(1, 2)                   # (B, T//k, C)
+        return self.norm_temporal(x)
+
+
+def downsample_key_padding_mask(mask: Optional[torch.Tensor], frames_per_patch: int) -> Optional[torch.Tensor]:
+    if mask is None:
+        return None
+    bsz, total_len = mask.shape
+    if total_len < frames_per_patch:
+        return mask.new_zeros((bsz, 0), dtype=mask.dtype)
+    out_len = total_len // frames_per_patch
+    trimmed = mask[:, : out_len * frames_per_patch]
+    grouped = trimmed.view(bsz, out_len, frames_per_patch)
+    return grouped.all(dim=-1)
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, num_heads: int, max_distance: int) -> None:
+        super().__init__()
+        if max_distance < 1:
+            raise ValueError("max_distance must be >= 1")
+        self.num_heads = num_heads
+        self.max_distance = max_distance
+        self.bias = nn.Parameter(torch.zeros(2 * max_distance - 1, num_heads))
+
+    def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        pos = torch.arange(seq_len, device=device)
+        rel = pos[:, None] - pos[None, :]
+        rel = rel.clamp(-self.max_distance + 1, self.max_distance - 1)
+        rel = rel + self.max_distance - 1
+        bias = self.bias[rel]
+        return bias.permute(2, 0, 1).to(dtype=dtype)
+
+
+class RelativeTransformerEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int,
+        dropout: float,
+        activation: str,
+    ) -> None:
+        super().__init__()
+        if d_model % nhead != 0:
+            raise ValueError("d_model must be divisible by nhead")
+        self.d_model = d_model
+        self.nhead = nhead
+        self.head_dim = d_model // nhead
+
+        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.resid_dropout = nn.Dropout(dropout)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.ff_dropout = nn.Dropout(dropout)
+
+        if activation == "gelu":
+            self.activation_fn = F.gelu
+        elif activation == "relu":
+            self.activation_fn = F.relu
+        else:
+            raise ValueError(f"unsupported activation: {activation}")
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        attn_bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = src.size()
+        qkv = self.qkv_proj(src)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(bsz, seq_len, self.nhead, self.head_dim)
+        k = k.view(bsz, seq_len, self.nhead, self.head_dim)
+        v = v.view(bsz, seq_len, self.nhead, self.head_dim)
+
+        attn_scores = torch.einsum("bthd,bshd->bhts", q, k) / math.sqrt(self.head_dim)
+
+        if src_mask is not None:
+            if src_mask.dtype == torch.bool:
+                attn_scores = attn_scores.masked_fill(src_mask.unsqueeze(0), float("-inf"))
+            else:
+                attn_scores = attn_scores + src_mask.unsqueeze(0)
+
+        if src_key_padding_mask is not None:
+            key_mask = src_key_padding_mask.unsqueeze(1).unsqueeze(2)
+            attn_scores = attn_scores.masked_fill(key_mask, float("-inf"))
+
+        if attn_bias is not None:
+            if attn_bias.dim() == 3:
+                attn_scores = attn_scores + attn_bias.unsqueeze(0)
+            elif attn_bias.dim() == 4:
+                attn_scores = attn_scores + attn_bias
+            else:
+                raise ValueError("attn_bias must be 3D or 4D tensor")
+
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        attn_weights = self.attn_dropout(attn_weights)
+        context = torch.einsum("bhts,bshd->bthd", attn_weights, v)
+        context = context.contiguous().view(bsz, seq_len, self.d_model)
+        attn_out = self.out_proj(context)
+        src = self.norm1(src + self.resid_dropout(attn_out))
+
+        ff_out = self.linear2(self.ff_dropout(self.activation_fn(self.linear1(src))))
+        src = self.norm2(src + self.resid_dropout(ff_out))
+        return src
+
+
+class RelativeTransformerEncoder(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int,
+        dropout: float,
+        activation: str,
+        relative_position_bias: Optional[RelativePositionBias],
+    ) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                RelativeTransformerEncoderLayer(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dim_feedforward=dim_feedforward,
+                    dropout=dropout,
+                    activation=activation,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(d_model)
+        self.rpb = relative_position_bias
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_key_padding_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        output = src
+        if self.rpb is not None:
+            attn_bias = self.rpb(src.size(1), device=src.device, dtype=src.dtype)
+        else:
+            attn_bias = None
+
+        for layer in self.layers:
+            output = layer(
+                output,
+                src_key_padding_mask=src_key_padding_mask,
+                attn_bias=attn_bias,
+            )
+
+        output = self.norm(output)
+        return output
+
+
+class ChordProjectionHead(nn.Module):
+    def __init__(self, d_model: int, vocab_sizes: Dict[str, int]) -> None:
+        super().__init__()
+        self.boundary_head = nn.Sequential(
+            nn.Linear(d_model, d_model // 2),
+            nn.GELU(),
+            nn.Linear(d_model // 2, 1),
+        )
+
+        projection_heads: Dict[str, nn.Module] = {}
+        for name, size in vocab_sizes.items():
+            projection_heads[name] = nn.Sequential(
+                nn.Linear(d_model, d_model // 2),
+                nn.GELU(),
+                nn.Linear(d_model // 2, size),
+            )
+        self.projection_heads = nn.ModuleDict(projection_heads)
+
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        boundary_logits = self.boundary_head(x).squeeze(-1)
+        outputs: Dict[str, torch.Tensor] = {"boundary": boundary_logits}
+        for comp, head in self.projection_heads.items():
+            outputs[comp] = head(x)
+        return outputs
+
+
+class KTokenDecoderLayer(nn.Module):
+    def __init__(self, d_model: int, nhead: int, mlp_ratio: int, dropout: float) -> None:
+        super().__init__()
+        if d_model % nhead != 0:
+            raise ValueError("d_model must be divisible by nhead")
+        self.d_model = d_model
+        self.nhead = nhead
+        self.head_dim = d_model // nhead
+
+        self.sa_qkv = nn.Linear(d_model, 3 * d_model)
+        self.sa_out = nn.Linear(d_model, d_model)
+        self.sa_ln = nn.LayerNorm(d_model)
+        self.sa_drop = nn.Dropout(dropout)
+
+        self.ca_q = nn.Linear(d_model, d_model)
+        self.ca_kv = nn.Linear(d_model, 2 * d_model)
+        self.ca_out = nn.Linear(d_model, d_model)
+        self.ca_ln = nn.LayerNorm(d_model)
+        self.ca_drop = nn.Dropout(dropout)
+
+        hidden = d_model * mlp_ratio
+        self.ff_ln = nn.LayerNorm(d_model)
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, d_model),
+        )
+        self.ff_drop = nn.Dropout(dropout)
+
+    def _attn(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        attn = torch.einsum("nlhd,nshd->nhls", q, k) / math.sqrt(q.size(-1))
+        attn = torch.softmax(attn, dim=-1)
+        ctx = torch.einsum("nhls,nshd->nlhd", attn, v)
+        ctx = ctx.contiguous().view(q.size(0), q.size(1), -1)
+        return ctx
+
+    def forward(self, x_tokens: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
+        bsz, length, k_tokens, d_model = x_tokens.shape
+        c_len = context.size(2)
+        batch_tokens = bsz * length
+
+        x = x_tokens.view(batch_tokens, k_tokens, d_model)
+        c = context.view(batch_tokens, c_len, d_model)
+
+        x_norm = self.sa_ln(x)
+        qkv = self.sa_qkv(x_norm)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(batch_tokens, k_tokens, self.nhead, self.head_dim)
+        k = k.view(batch_tokens, k_tokens, self.nhead, self.head_dim)
+        v = v.view(batch_tokens, k_tokens, self.nhead, self.head_dim)
+        sa_ctx = self._attn(q, k, v)
+        x = x + self.sa_drop(self.sa_out(sa_ctx))
+
+        x_norm = self.ca_ln(x)
+        q = self.ca_q(x_norm).view(batch_tokens, k_tokens, self.nhead, self.head_dim)
+        kv = self.ca_kv(c)
+        k, v = kv.chunk(2, dim=-1)
+        k = k.view(batch_tokens, c_len, self.nhead, self.head_dim)
+        v = v.view(batch_tokens, c_len, self.nhead, self.head_dim)
+        ca_ctx = self._attn(q, k, v)
+        x = x + self.ca_drop(self.ca_out(ca_ctx))
+
+        x_norm = self.ff_ln(x)
+        x = x + self.ff_drop(self.ff(x_norm))
+
+        return x.view(bsz, length, k_tokens, d_model)
+
+
+def build_rpb(config: Dict[str, Any]) -> RelativePositionBias:
+    return RelativePositionBias(
+        num_heads=config["n_head"],
+        max_distance=config["n_beats"] * config["label_resolution"],
+    )
+
+
+def build_encoder(config: Dict[str, Any]) -> RelativeTransformerEncoder:
+    rpb = build_rpb(config)
+    return RelativeTransformerEncoder(
+        num_layers=config["num_encoder_layers"],
+        d_model=config["d_model"],
+        nhead=config["n_head"],
+        dim_feedforward=config["dim_feedforward"],
+        dropout=config["dropout"],
+        activation="gelu",
+        relative_position_bias=rpb,
+    )
+
+
+def build_patch_embedding(config: Dict[str, Any]) -> PatchEmbedding:
+    return PatchEmbedding(
+        d_model=config["d_model"],
+        frames_per_patch=config["frames_per_patch"],
+        expansion=2,
+    )
+
+

models/model.py

@@ -0,0 +1,686 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Dict, Any, Optional
+
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, d_model: int, frames_per_patch: int = 6, expansion: int = 2):
+        super().__init__()
+        self.d_model = d_model
+        self.frames_per_patch = frames_per_patch
+        # Frame embedding (collapse pitch dim)
+        self.conv2d = nn.Conv2d(
+            in_channels=1,
+            out_channels=d_model,
+            kernel_size=(88, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+        )
+        self.norm_frame = nn.LayerNorm(d_model)
+        # anti-aliasing conv on time axis
+        self.aa = nn.Conv1d(d_model, d_model, kernel_size=3, stride=1,
+                            padding=1, groups=d_model, bias=False)
+        
+        # Late temporal pooling (downsample frames -> patches)
+        self.glu_conv = nn.Conv1d(
+            in_channels=d_model,
+            out_channels=d_model * expansion * 2,
+            kernel_size=frames_per_patch,
+            stride=frames_per_patch,
+            padding=0,
+            bias=True,
+        )
+        self.project = nn.Conv1d(
+            in_channels=d_model * expansion,
+            out_channels=d_model,
+            kernel_size=1,
+        )
+        self.norm_temporal = nn.LayerNorm(d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, 1, 88, T)
+        x = self.conv2d(x)                      # (B, C, 1, T)
+        x = x.squeeze(2).transpose(1, 2)        # (B, T, C)
+        x = self.norm_frame(x)
+
+        # anti-aliased and temporal pooling
+        x = x.transpose(1, 2)                   # (B, C, T)
+        x = self.aa(x)                          # (B, C, T)
+        v, g = self.glu_conv(x).chunk(2, dim=1)
+        x = self.project(v * torch.sigmoid(g))  # (B, C, T//k)
+        x = x.transpose(1, 2)                   # (B, T//k, C)
+        return self.norm_temporal(x)
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x.transpose(0, 1)
+        x = x + self.pe[:x.size(0), :]
+        x = self.dropout(x)
+        return x.transpose(0, 1)
+
+
+class RelativeTransformerEncoderLayer(nn.Module):
+    def __init__(self, d_model: int, nhead: int, dim_feedforward: int, dropout: float, activation: str = 'gelu'):
+        super().__init__()
+        if d_model % nhead != 0:
+            raise ValueError("d_model must be divisible by nhead.")
+        self.d_model = d_model
+        self.nhead = nhead
+        self.head_dim = d_model // nhead
+
+        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.resid_dropout = nn.Dropout(dropout)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.ff_dropout = nn.Dropout(dropout)
+
+        if activation == 'gelu':
+            self.activation_fn = F.gelu
+        elif activation == 'relu':
+            self.activation_fn = F.relu
+        else:
+            raise ValueError(f"Unsupported activation: {activation}")
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        attn_bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # src: (B, T_new, C)
+        bsz, seq_len_new, _ = src.size()
+
+        qkv = self.qkv_proj(src)  # (B, T_new, 3*C)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(bsz, seq_len_new, self.nhead, self.head_dim)
+        k_all = k.view(bsz, seq_len_new, self.nhead, self.head_dim)
+        v_all = v.view(bsz, seq_len_new, self.nhead, self.head_dim)
+
+        # Attention: queries are only for current tokens; keys include past+current
+        attn_scores = torch.einsum('bthd,bshd->bhts', q, k_all) / math.sqrt(self.head_dim)  # (B, H, T_new, T_total)
+
+        # Additive or boolean mask over attention logits
+        if src_mask is not None:
+            if src_mask.dtype == torch.bool:
+                attn_scores = attn_scores.masked_fill(src_mask.unsqueeze(0), float('-inf'))
+            else:
+                attn_scores = attn_scores + src_mask.unsqueeze(0)
+
+        # Key padding mask
+        if src_key_padding_mask is not None:
+            key_mask = src_key_padding_mask.unsqueeze(1).unsqueeze(2)  # (B,1,1,T)
+            attn_scores = attn_scores.masked_fill(key_mask, float('-inf'))
+
+        if attn_bias is not None:
+            # Support 3D (H, T, T) or 4D (B, H, T, T)
+            if attn_bias.dim() == 3:
+                attn_scores = attn_scores + attn_bias.unsqueeze(0)
+            elif attn_bias.dim() == 4:
+                attn_scores = attn_scores + attn_bias
+            else:
+                raise ValueError("attn_bias must be 3D or 4D tensor if provided")
+
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        context = torch.einsum('bhts,bshd->bthd', attn_weights, v_all)  # (B, T_new, H, D)
+        context = context.contiguous().view(bsz, seq_len_new, self.d_model)  # (B, T_new, C)
+        attn_out = self.out_proj(context)
+
+        src = src + self.resid_dropout(attn_out)
+        src = self.norm1(src)
+
+        ff = self.linear2(self.ff_dropout(self.activation_fn(self.linear1(src))))
+        src = src + self.resid_dropout(ff)
+        src = self.norm2(src)
+
+        return src
+
+def downsample_key_padding_mask(mask: torch.Tensor, frames_per_patch: int) -> torch.Tensor:
+    # mask: (B, T) where True denotes padding.
+    bsz, total_len = mask.shape
+    if total_len < frames_per_patch:
+        # No valid output tokens from temporal pooling
+        return mask.new_ones((bsz, 0), dtype=mask.dtype)
+    out_len = total_len // frames_per_patch
+    trimmed = mask[:, :out_len * frames_per_patch]
+    grouped = trimmed.view(bsz, out_len, frames_per_patch)
+    return grouped.all(dim=-1)
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, num_heads: int, max_distance: int):
+        super().__init__()
+        if max_distance < 1:
+            raise ValueError("max_distance must be >= 1")
+        self.max_distance = max_distance
+        self.num_heads = num_heads
+        # Table over relative distances in [-max_distance+1, max_distance-1]
+        self.bias = nn.Parameter(torch.zeros(2 * max_distance - 1, num_heads))
+
+    def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        # Compute clipped relative position indices
+        pos = torch.arange(seq_len, device=device)
+        rel = pos[:, None] - pos[None, :]  # (T, T)
+        rel = rel.clamp(-self.max_distance + 1, self.max_distance - 1)
+        rel = rel + self.max_distance - 1  # shift to [0, 2*max_distance-2]
+        bias = self.bias[rel]  # (T, T, H)
+        return bias.permute(2, 0, 1).to(dtype=dtype)  # (H, T, T)
+
+
+class RelativeTransformerEncoder(nn.Module):
+    def __init__(self, num_layers: int, d_model: int, nhead: int, dim_feedforward: int, 
+                 dropout: float, activation: str = 'gelu', relative_position_bias = None):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            RelativeTransformerEncoderLayer(
+                d_model=d_model,
+                nhead=nhead,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+            ) for _ in range(num_layers)
+        ])
+        self.norm = nn.LayerNorm(d_model)
+        self.rpb = relative_position_bias
+        self.nhead = nhead
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        output = src
+
+        if self.rpb is not None:
+            attn_bias = self.rpb(src.size(1), device=src.device, dtype=src.dtype)
+        else:
+            attn_bias = None
+
+        for mod in self.layers:
+            output = mod(
+                output,
+                src_key_padding_mask=src_key_padding_mask,
+                attn_bias=attn_bias,
+            )
+
+        output = self.norm(output)
+        return output
+
+
+class ChordDecomposeProjection(nn.Module):
+    def __init__(self, d_model: int, vocab_sizes: Dict[str, int]):
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_sizes = vocab_sizes
+        self.boundary_head = nn.Sequential(
+            nn.Linear(d_model, d_model // 2),
+            nn.GELU(),
+            nn.Linear(d_model // 2, 1),
+        )
+
+        self.projection_heads = nn.ModuleDict()
+        for comp, size in self.vocab_sizes.items():
+            self.projection_heads[comp] = nn.Sequential(
+                nn.Linear(d_model, d_model // 2),
+                nn.GELU(),
+                nn.Linear(d_model // 2, size),
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        boundary_logits = self.boundary_head(x)
+
+        output = {'boundary': boundary_logits.squeeze(-1)}
+        for comp, head in self.projection_heads.items():
+            output[comp] = head(x)
+
+        return output
+
+
+class ChordRecognitionModel(nn.Module):
+    def __init__(self, model_config: Dict[str, Any], vocab_sizes: Dict[str, int]):
+        super().__init__()
+        self.config = model_config
+        self.vocab_sizes = vocab_sizes
+        self.d_model = self.config['d_model']
+        # Encoder: shared patch embedding and relative transformer (unchanged)
+        self.embedding = PatchEmbedding(
+            d_model=self.d_model,
+            frames_per_patch=self.config['frames_per_patch'],
+            expansion=2,
+        )
+
+        self.input_dropout = nn.Dropout(self.config['dropout'])
+
+        rpb = RelativePositionBias(
+            num_heads=self.config['n_head'],
+            max_distance=self.config['n_beats'] * self.config['label_resolution']
+        )
+        self.relative_transformer_encoder = RelativeTransformerEncoder(
+            num_layers=self.config['num_encoder_layers'],
+            d_model=self.d_model,
+            nhead=self.config['n_head'],
+            dim_feedforward=self.config['dim_feedforward'],
+            dropout=self.config['dropout'],
+            activation='gelu',
+            relative_position_bias=rpb,
+        )
+
+        # Boundary head, smoother, and FiLM gating
+        d_b = max(1, self.d_model // 4)
+        k_b = int(self.config.get('boundary_kernel', 5))
+        self.boundary_head = nn.Linear(self.d_model, 1)
+        self.boundary_smoother = nn.Conv1d(
+            in_channels=1,
+            out_channels=1,
+            kernel_size=k_b,
+            padding=k_b // 2,
+            groups=1,
+            bias=True,
+        )
+        self.boundary_e0 = nn.Parameter(torch.zeros(d_b))
+        self.boundary_e1 = nn.Parameter(torch.randn(d_b) * 0.02)
+        # Optional key context (3-part setting). Infer size from vocab_sizes if provided
+        self.Vq = int(vocab_sizes.get('quality', 0))
+        self.Vr = int(vocab_sizes.get('root', 0))
+        self.Vb = int(vocab_sizes.get('bass', 0))
+
+        # FiLM layers take boundary embedding
+        self.film_ln_in = nn.LayerNorm(self.d_model + d_b)
+        self.film_ln_h = nn.LayerNorm(self.d_model)
+        self.film_mlp = nn.Linear(self.d_model + d_b, 2 * self.d_model)
+
+        # Triple-token decoder: embeddings and heads
+        self.mask_id_q = int(self.config.get('mask_id_q', self.Vq))
+        self.mask_id_r = int(self.config.get('mask_id_r', self.Vr))
+        self.mask_id_b = int(self.config.get('mask_id_b', self.Vb))
+        self.emb_q = nn.Embedding(self.Vq + 1, self.d_model)
+        self.emb_r = nn.Embedding(self.Vr + 1, self.d_model)
+        self.emb_b = nn.Embedding(self.Vb + 1, self.d_model)
+
+        dec_heads = int(self.config.get('dec_heads', 4))
+        dec_mlp_ratio = int(self.config.get('dec_mlp_ratio', 4))
+        dec_layers = int(self.config.get('dec_layers', 1))
+        dec_dropout = float(self.config.get('dec_dropout', 0.1))
+        self.window_radius = int(self.config.get('window_radius', 2))
+        self.decoder_layers = nn.ModuleList([
+            KTokenDecoderLayer(
+                d_model=self.d_model,
+                nhead=dec_heads,
+                mlp_ratio=dec_mlp_ratio,
+                dropout=dec_dropout,
+            ) for _ in range(dec_layers)
+        ])
+        self.dec_norm = nn.LayerNorm(self.d_model)
+        self.head_q = nn.Linear(self.d_model, self.Vq)
+        self.head_r = nn.Linear(self.d_model, self.Vr)
+        self.head_b = nn.Linear(self.d_model, self.Vb)
+        
+        # Legacy decompose projection for compatibility when training in decompose mode
+        self.chord_decompose_projection = ChordDecomposeProjection(self.d_model, self.vocab_sizes)
+
+    def forward(self, encoder_input: torch.Tensor, 
+                src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        H, _ = self._encode(encoder_input, src_key_padding_mask)
+        out = self.chord_decompose_projection(H)
+        return out
+
+    # ===== Utilities =====
+    def _encode(self, encoder_input: torch.Tensor, 
+                src_key_padding_mask: Optional[torch.Tensor]) -> (torch.Tensor, torch.Tensor):
+        x = self.embedding(encoder_input)
+        mask_down = None
+        if src_key_padding_mask is not None:
+            mask_down = downsample_key_padding_mask(src_key_padding_mask, self.config['frames_per_patch'])
+        x = self.input_dropout(x)
+        H = self.relative_transformer_encoder(x, src_key_padding_mask=mask_down)
+        boundary_logits = self.boundary_head(H).squeeze(-1)
+        return H, boundary_logits
+
+    def _smooth_boundary(self, boundary_logits: torch.Tensor) -> torch.Tensor:
+        b = boundary_logits.unsqueeze(1)
+        smoothed = self.boundary_smoother(b)
+        return torch.sigmoid(smoothed.squeeze(1))
+
+    def _apply_film(self, H: torch.Tensor, b_soft: torch.Tensor) -> torch.Tensor:
+        B, T, D = H.shape
+        e0 = self.boundary_e0.view(1, 1, -1).expand(B, T, -1)
+        e1 = self.boundary_e1.view(1, 1, -1).expand(B, T, -1)
+        # soft embedding for boundary
+        eb = b_soft.unsqueeze(-1) * e1 + (1.0 - b_soft).unsqueeze(-1) * e0
+        film_in = torch.cat([H, eb], dim=-1)
+
+        # layer norm and linear projection
+        film_in = self.film_ln_in(film_in)
+        gamma, beta = self.film_mlp(film_in).chunk(2, dim=-1)
+        Z = self.film_ln_h(H) * (1.0 + gamma) + beta
+        return Z
+
+    def _build_local_windows(self, H: torch.Tensor, radius: int) -> torch.Tensor:
+        x = H.transpose(1, 2)
+        padded = F.pad(x, (radius, radius), mode='replicate')
+        win = padded.unfold(dimension=2, size=2 * radius + 1, step=1)
+        win = win.permute(0, 2, 3, 1).contiguous()
+        return win
+
+    def _build_context(self, H: torch.Tensor, Z: torch.Tensor, b_soft: torch.Tensor) -> torch.Tensor:
+        local = self._build_local_windows(H, self.window_radius)
+        z = Z.unsqueeze(2)
+        parts = [z, local]
+        C = torch.cat(parts, dim=2)
+        return C
+
+    def _embed_tokens(self, ids_q: torch.Tensor, ids_r: torch.Tensor, ids_b: torch.Tensor) -> torch.Tensor:
+        xq = self.emb_q(ids_q)
+        xr = self.emb_r(ids_r)
+        xb = self.emb_b(ids_b)
+        X = torch.stack([xq, xr, xb], dim=2)
+        return X
+
+    def _run_decoder(self, X: torch.Tensor, C: torch.Tensor):
+        x = X
+        for layer in self.decoder_layers:
+            x = layer(x, C)
+        x = self.dec_norm(x)
+        xq = x[:, :, 0, :]
+        xr = x[:, :, 1, :]
+        xb = x[:, :, 2, :]
+        logits_q = self.head_q(xq)
+        logits_r = self.head_r(xr)
+        logits_b = self.head_b(xb)
+        return logits_q, logits_r, logits_b
+
+    # ===== Training forward =====
+    def forward_train(self, encoder_input: torch.Tensor,
+                      targets: Dict[str, torch.Tensor],
+                      src_key_padding_mask: Optional[torch.Tensor] = None,
+                      target_mask: Optional[torch.Tensor] = None) -> Dict[str, Any]:
+        device = encoder_input.device
+        H, boundary_logits = self._encode(encoder_input, src_key_padding_mask)
+        # No key prediction/context in training
+        b_soft = self._smooth_boundary(boundary_logits)
+        Z = self._apply_film(H, b_soft)
+        C = self._build_context(H, Z, b_soft)
+
+        tgt_q = targets['quality']
+        tgt_r = targets['root']
+        tgt_b = targets['bass']
+        B, T = tgt_q.shape
+        if target_mask is None:
+            target_mask = torch.ones(B, T, dtype=torch.bool, device=device)
+
+        # mask n slots randomly per (B,T) across 3 slots [q,r,b]
+        k_rand = torch.randint(1, 4, (B, T), device=device)
+        rand_scores = torch.rand(B, T, 3, device=device)
+        top_vals, top_idx = torch.topk(rand_scores, k=3, dim=-1)
+        mask_slots = torch.zeros(B, T, 3, dtype=torch.bool, device=device)
+        # enable first k indices per position
+        for kk in range(1, 4):
+            sel = (k_rand == kk)
+            if sel.any():
+                idx_sel = top_idx[sel][:, :kk]
+                row = mask_slots[sel]
+                if idx_sel.numel() > 0:
+                    row.scatter_(dim=1, index=idx_sel, value=True)
+                mask_slots[sel] = row
+
+        ids_q = tgt_q.clone()
+        ids_r = tgt_r.clone()
+        ids_b = tgt_b.clone()
+        ids_q[mask_slots[:, :, 0]] = self.mask_id_q
+        ids_r[mask_slots[:, :, 1]] = self.mask_id_r
+        ids_b[mask_slots[:, :, 2]] = self.mask_id_b
+
+        X = self._embed_tokens(ids_q, ids_r, ids_b)
+        logits_q, logits_r, logits_b = self._run_decoder(X, C)
+
+        def ce_masked(logits: Optional[torch.Tensor], target: torch.Tensor, slot_mask: torch.Tensor) -> torch.Tensor:
+            # Build supervision mask and safe targets to avoid CUDA asserts from out-of-range labels
+            m = slot_mask & target_mask  # (B,T) supervised locations
+            num_classes = logits.size(-1)
+            safe_target = torch.where(
+                m,
+                target.clamp(min=0, max=num_classes - 1),
+                torch.zeros_like(target)
+            )
+            ce = F.cross_entropy(logits.transpose(1, 2), safe_target, reduction='none')
+            denom = m.float().sum().clamp(min=1.0)
+            return (ce * m.float()).sum() / denom
+
+        loss_q = ce_masked(logits_q, tgt_q, mask_slots[:, :, 0])
+        loss_r = ce_masked(logits_r, tgt_r, mask_slots[:, :, 1])
+        loss_b = ce_masked(logits_b, tgt_b, mask_slots[:, :, 2])
+
+        bce = F.binary_cross_entropy_with_logits(boundary_logits, 
+                                                 targets['boundary'].to(boundary_logits.dtype), 
+                                                 pos_weight=torch.tensor(2.0, device=device),
+                                                 reduction='none')
+        loss_boundary = (bce * target_mask.float()).sum() / target_mask.float().sum().clamp(min=1.0)
+        total_loss = loss_q + loss_r + loss_b + loss_boundary * 3
+
+        with torch.no_grad():
+            stats = {}
+            for name, logits, target, m in [
+                ('quality', logits_q, tgt_q, mask_slots[:, :, 0]),
+                ('root', logits_r, tgt_r, mask_slots[:, :, 1]),
+                ('bass', logits_b, tgt_b, mask_slots[:, :, 2]),
+            ]:
+                if logits is None:
+                    stats[f'acc_{name}'] = 0.0
+                    stats[f'conf_{name}'] = 0.0
+                    stats[f'ece_{name}'] = 0.0
+                else:
+                    pred = logits.argmax(dim=-1)
+                    sel = (m & target_mask)
+                    denom = sel.float().sum().clamp(min=1.0)
+                    acc = (pred[sel] == target[sel]).float().sum() / denom
+                    prob = logits.float().softmax(dim=-1)
+                    conf = prob.max(dim=-1).values
+                    mean_conf = conf[sel].sum() / denom
+                    # simple ECE
+                    ece = torch.tensor(0.0, device=device)
+                    bins = torch.linspace(0, 1, steps=11, device=device)
+                    conf_flat = conf[sel]
+                    pred_flat = pred[sel]
+                    tgt_flat = target[sel]
+                    for i in range(10):
+                        lo, hi = bins[i], bins[i+1]
+                        mask_bin = (conf_flat >= lo) & (conf_flat < hi if i < 9 else conf_flat <= hi)
+                        if mask_bin.sum() > 0:
+                            acc_bin = (pred_flat[mask_bin] == tgt_flat[mask_bin]).float().mean()
+                            conf_bin = conf_flat[mask_bin].mean()
+                            ece = ece + (mask_bin.float().mean() * (acc_bin - conf_bin).abs())
+                    stats[f'acc_{name}'] = acc.item()
+                    stats[f'conf_{name}'] = mean_conf.item()
+                    stats[f'ece_{name}'] = ece.item()
+
+        return {
+            'loss': total_loss,
+            'loss_map': {
+                'quality': loss_q,
+                'root': loss_r,
+                'bass': loss_b,
+                'boundary': loss_boundary,
+            },
+            'logits': {
+                'quality': logits_q,
+                'root': logits_r,
+                'bass': logits_b,
+            },
+            'mask_slots': mask_slots,  # (B,T,3) bool in order [q,r,b]
+            'boundary_logits': boundary_logits,
+            'stats': stats,
+        }
+
+    # ===== Inference forward =====
+    def forward_infer(self, encoder_input: torch.Tensor,
+                      src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        device = encoder_input.device
+        H, boundary_logits = self._encode(encoder_input, src_key_padding_mask)
+        # No key prediction/context in inference
+        b_soft = self._smooth_boundary(boundary_logits)
+        Z = self._apply_film(H, b_soft)
+        C = self._build_context(H, Z, b_soft)
+
+        B, T, _ = H.shape
+        ids_q = torch.full((B, T), self.mask_id_q, dtype=torch.long, device=device)
+        ids_r = torch.full((B, T), self.mask_id_r, dtype=torch.long, device=device)
+        ids_b = torch.full((B, T), self.mask_id_b, dtype=torch.long, device=device)
+        filled_q = torch.zeros((B, T), dtype=torch.bool, device=device)
+        filled_r = torch.zeros((B, T), dtype=torch.bool, device=device)
+        filled_b = torch.zeros((B, T), dtype=torch.bool, device=device)
+
+        # Track decode order per time step: 0=quality, 1=root, 2=bass
+        decode_order = torch.full((B, T, 3), -1, dtype=torch.long, device=device)
+        order_pos = 0
+
+        for step in (3, 2, 1):
+            X = self._embed_tokens(ids_q, ids_r, ids_b)
+            logits_q, logits_r, logits_b = self._run_decoder(X, C)
+            pq = logits_q.softmax(dim=-1)
+            pr = logits_r.softmax(dim=-1)
+            pb = logits_b.softmax(dim=-1)
+            conf_q = pq.max(dim=-1).values
+            conf_r = pr.max(dim=-1).values
+            conf_b = pb.max(dim=-1).values
+            conf_q = conf_q.masked_fill(filled_q, float('-inf'))
+            conf_r = conf_r.masked_fill(filled_r, float('-inf'))
+            conf_b = conf_b.masked_fill(filled_b, float('-inf'))
+            conf = torch.stack([conf_q, conf_r, conf_b], dim=-1)
+            take_slot = conf.argmax(dim=-1)
+
+            # record order at this step
+            decode_order[:, :, order_pos] = take_slot
+            order_pos += 1
+            
+            pred_q = logits_q.argmax(dim=-1)
+            commit_q = (take_slot == 0) | ((step == 1) & (~filled_q))
+            ids_q[commit_q] = pred_q[commit_q]
+            filled_q = filled_q | commit_q
+            
+            pred_r = logits_r.argmax(dim=-1)
+            commit_r = (take_slot == 1) | ((step == 1) & (~filled_r))
+            ids_r[commit_r] = pred_r[commit_r]
+            filled_r = filled_r | commit_r
+            
+            pred_b = logits_b.argmax(dim=-1)
+            commit_b = (take_slot == 2) | ((step == 1) & (~filled_b))
+            ids_b[commit_b] = pred_b[commit_b]
+            filled_b = filled_b | commit_b
+
+        # final confidences
+        X = self._embed_tokens(ids_q, ids_r, ids_b)
+        logits_q, logits_r, logits_b = self._run_decoder(X, C)
+        conf_q = logits_q.softmax(dim=-1).max(dim=-1).values
+        conf_r = logits_r.softmax(dim=-1).max(dim=-1).values
+        conf_b = logits_b.softmax(dim=-1).max(dim=-1).values
+        
+        return {
+            'quality': ids_q,
+            'root': ids_r,
+            'bass': ids_b,
+            'conf_quality': conf_q,
+            'conf_root': conf_r,
+            'conf_bass': conf_b,
+            'boundary': boundary_logits,
+            'decode_order': decode_order,
+        }
+
+
+class KTokenDecoderLayer(nn.Module):
+    def __init__(self, d_model: int, nhead: int, mlp_ratio: int, dropout: float):
+        super().__init__()
+        self.d_model = d_model
+        self.nhead = nhead
+        self.head_dim = d_model // nhead
+        if self.head_dim * nhead != d_model:
+            raise ValueError("d_model must be divisible by nhead")
+
+        # self-attention over K tokens
+        self.sa_qkv = nn.Linear(d_model, 3 * d_model)
+        self.sa_out = nn.Linear(d_model, d_model)
+        self.sa_ln = nn.LayerNorm(d_model)
+        self.sa_drop = nn.Dropout(dropout)
+
+        # cross-attention to context
+        self.ca_q = nn.Linear(d_model, d_model)
+        self.ca_kv = nn.Linear(d_model, 2 * d_model)
+        self.ca_out = nn.Linear(d_model, d_model)
+        self.ca_ln = nn.LayerNorm(d_model)
+        self.ca_drop = nn.Dropout(dropout)
+
+        # ffn
+        hidden = d_model * mlp_ratio
+        self.ff_ln = nn.LayerNorm(d_model)
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, d_model),
+        )
+        self.ff_drop = nn.Dropout(dropout)
+
+    def _attn(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        # q,k,v: (N, L, H, D)
+        attn = torch.einsum('nlhd,nshd->nhls', q, k) / math.sqrt(q.size(-1))  # (N,H,L,S)
+        attn = torch.softmax(attn, dim=-1)
+        ctx = torch.einsum('nhls,nshd->nlhd', attn, v)  # (N,L,H,D)
+        ctx = ctx.contiguous().view(q.size(0), q.size(1), -1)  # (N,L,C)
+        return ctx
+
+    def forward(self, X: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
+        # X: (B,T,K,D), C: (B,T,Lc,D)
+        B, T, K, D = X.shape
+        Lc = C.size(2)
+        N = B * T
+        # reshape
+        x = X.view(N, K, D)
+        c = C.view(N, Lc, D)
+
+        # self-attn (over K)
+        x_norm = self.sa_ln(x)
+        qkv = self.sa_qkv(x_norm)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(N, K, self.nhead, self.head_dim)
+        k = k.view(N, K, self.nhead, self.head_dim)
+        v = v.view(N, K, self.nhead, self.head_dim)
+        sa_ctx = self._attn(q, k, v)  # (N,K,C)
+        x = x + self.sa_drop(self.sa_out(sa_ctx))
+
+        # cross-attn (queries = tokens, keys/values = context)
+        x_norm = self.ca_ln(x)
+        q = self.ca_q(x_norm).view(N, K, self.nhead, self.head_dim)
+        kv = self.ca_kv(c)
+        k, v = kv.chunk(2, dim=-1)
+        k = k.view(N, Lc, self.nhead, self.head_dim)
+        v = v.view(N, Lc, self.nhead, self.head_dim)
+        ca_ctx = self._attn(q, k, v)  # (N,K,C)
+        x = x + self.ca_drop(self.ca_out(ca_ctx))
+
+        # ffn
+        x_norm = self.ff_ln(x)
+        x = x + self.ff_drop(self.ff(x_norm))
+
+        return x.view(B, T, K, D)

models/variants.py

@@ -0,0 +1,654 @@
+import math
+from typing import Dict, Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .components import (
+    downsample_key_padding_mask,
+    KTokenDecoderLayer,
+    ChordProjectionHead,
+    build_encoder,
+    build_patch_embedding,
+)
+
+
+def _to_patch_input(x_bt88: torch.Tensor) -> torch.Tensor:
+    # Accepts (B, T, 88) -> returns (B, 1, 88, T)
+    return x_bt88.transpose(1, 2).unsqueeze(1).contiguous()
+
+class BaseEncoder(nn.Module):
+    def __init__(self, model_config: Dict[str, Any]):
+        super().__init__()
+        self.config = model_config
+        self.d_model = self.config["d_model"]
+        self.embedding = build_patch_embedding(self.config)
+        self.input_dropout = nn.Dropout(self.config["dropout"])
+        self.encoder = build_encoder(self.config)
+        self.boundary_head = nn.Linear(self.d_model, 1)
+
+    def encode(self, encoder_input_bt88: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        x_pr = _to_patch_input(encoder_input_bt88)
+        x = self.embedding(x_pr)
+        mask_down = None
+        if src_key_padding_mask is not None:
+            mask_down = downsample_key_padding_mask(src_key_padding_mask, self.config["frames_per_patch"])
+        x = self.input_dropout(x)
+        h = self.encoder(x, src_key_padding_mask=mask_down)
+        boundary_logits = self.boundary_head(h).squeeze(-1)
+        return h, boundary_logits
+
+
+class BaselineLinearModel(nn.Module):
+    """PatchEmbedding + TransformerEncoder + linear heads (baseline)."""
+
+    def __init__(self, model_config: Dict[str, Any], vocab_sizes: Dict[str, int]):
+        super().__init__()
+        self.config = model_config
+        self.encoder = BaseEncoder(model_config)
+        self.proj = ChordProjectionHead(model_config["d_model"], vocab_sizes)
+        self.use_key = ('key' in vocab_sizes)
+
+    def forward_train(
+        self,
+        encoder_input: torch.Tensor,
+        targets: Dict[str, torch.Tensor],
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        target_mask: Optional[torch.Tensor] = None,
+        vocabs: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        device = encoder_input.device
+        h, boundary_logits = self.encoder.encode(encoder_input, src_key_padding_mask)
+        outputs = self.proj(h)
+
+        bsz, t_len = targets["quality"].shape
+        if target_mask is None:
+            target_mask = torch.ones(bsz, t_len, dtype=torch.bool, device=device)
+
+        def ce_masked(logits: torch.Tensor, target: torch.Tensor, mask: torch.Tensor, comp_name: str) -> torch.Tensor:
+            num_classes = logits.size(-1)
+            comp_pad = vocabs.get(f"{comp_name}_pad_idx", vocabs["pad_idx"]) if vocabs is not None else 0
+            valid = mask & (target != comp_pad)
+            safe_target = torch.where(valid, target.clamp(min=0, max=num_classes - 1), torch.zeros_like(target))
+            ce = F.cross_entropy(logits.transpose(1, 2), safe_target, reduction="none")
+            denom = valid.float().sum().clamp(min=1.0)
+            return (ce * valid.float()).sum() / denom
+
+        loss_q = ce_masked(outputs["quality"], targets["quality"], target_mask, "quality")
+        loss_r = ce_masked(outputs["root"], targets["root"], target_mask, "root")
+        loss_b = ce_masked(outputs["bass"], targets["bass"], target_mask, "bass")
+        if self.use_key and ("key" in outputs) and ("key" in targets):
+            loss_k = ce_masked(outputs["key"], targets["key"], target_mask, "key")
+        else:
+            loss_k = torch.tensor(0.0, device=device)
+        bce = F.binary_cross_entropy_with_logits(boundary_logits, targets["boundary"].to(boundary_logits.dtype), reduction="none")
+        loss_boundary = (bce * target_mask.float()).sum() / target_mask.float().sum().clamp(min=1.0)
+        total_loss = loss_q + loss_r + loss_b + loss_k + loss_boundary * 3.0
+
+        with torch.no_grad():
+            logits = {k: v for k, v in outputs.items() if k in ("quality", "root", "bass")}
+        return {
+            "loss": total_loss,
+            "loss_map": {
+                "quality": loss_q,
+                "root": loss_r,
+                "bass": loss_b,
+                "key": loss_k,
+                "boundary": loss_boundary,
+            },
+            "logits": logits,
+            "boundary_logits": boundary_logits,
+        }
+
+    def forward_infer(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        h, boundary_logits = self.encoder.encode(encoder_input, src_key_padding_mask)
+        outputs = self.proj(h)
+        comps = ["quality","root","bass"] + (["key"] if self.use_key and ("key" in outputs) else [])
+        ids = {k: outputs[k].argmax(dim=-1) for k in comps}
+        conf = {f"conf_{k}": outputs[k].softmax(dim=-1).max(dim=-1).values for k in comps}
+        ids.update(conf)
+        ids["boundary"] = boundary_logits
+        return ids
+
+
+class FiLMContextLinearModel(nn.Module):
+    """FiLM injection + local context window + per-component context projector (linear heads)."""
+
+    def __init__(self, model_config: Dict[str, Any], vocab_sizes: Dict[str, int]):
+        super().__init__()
+        self.config = model_config
+        self.encoder = BaseEncoder(model_config)
+        d_model = model_config["d_model"]
+        self.window_radius = int(model_config["window_radius"])  # Lc = 2R+1 + 1(Z)
+
+        # FiLM parameters (copy of CR-model style)
+        self.d_b = max(1, d_model // 4)
+        k_b = int(model_config["boundary_kernel"])
+        self.boundary_smoother = nn.Conv1d(1, 1, kernel_size=k_b, padding=k_b // 2, bias=True)
+        self.boundary_e0 = nn.Parameter(torch.zeros(self.d_b))
+        self.boundary_e1 = nn.Parameter(torch.randn(self.d_b) * 0.02)
+
+        self.film_ln_in = nn.LayerNorm(d_model + self.d_b)
+        self.film_ln_h = nn.LayerNorm(d_model)
+        self.film_mlp = nn.Linear(d_model + self.d_b, 2 * d_model)
+
+        # Per-component context heads
+        self.comp_names = ["quality", "root", "bass"]
+        if 'key' in vocab_sizes:
+            self.comp_names.append('key')
+        self.vocab_sizes = vocab_sizes
+        self.attn_mlp = nn.ModuleDict()
+        self.comp_proj = nn.ModuleDict()
+        for comp in self.comp_names:
+            self.attn_mlp[comp] = nn.Sequential(
+                nn.Linear(model_config["d_model"], model_config["d_model"] // 2),
+                nn.GELU(),
+                nn.Linear(model_config["d_model"] // 2, 1),  # score per context token
+            )
+            self.comp_proj[comp] = nn.Sequential(
+                nn.Linear(model_config["d_model"], model_config["d_model"]),
+                nn.GELU(),
+                nn.Linear(model_config["d_model"], self.vocab_sizes[comp]),
+            )
+
+    def _smooth_boundary(self, boundary_logits: torch.Tensor) -> torch.Tensor:
+        b = boundary_logits.unsqueeze(1)
+        smoothed = self.boundary_smoother(b)
+        return torch.sigmoid(smoothed.squeeze(1))
+
+    def _apply_film(self, h: torch.Tensor, b_soft: torch.Tensor) -> torch.Tensor:
+        bsz, t_len, d_model = h.shape
+        e0 = self.boundary_e0.view(1, 1, -1).expand(bsz, t_len, -1)
+        e1 = self.boundary_e1.view(1, 1, -1).expand(bsz, t_len, -1)
+        eb = b_soft.unsqueeze(-1) * e1 + (1.0 - b_soft).unsqueeze(-1) * e0
+        film_in = torch.cat([h, eb], dim=-1)
+        film_in = self.film_ln_in(film_in)
+        gamma, beta = self.film_mlp(film_in).chunk(2, dim=-1)
+        z = self.film_ln_h(h) * (1.0 + gamma) + beta
+        return z
+
+    def _build_local_windows(self, h: torch.Tensor, radius: int) -> torch.Tensor:
+        x = h.transpose(1, 2)
+        padded = F.pad(x, (radius, radius), mode="replicate")
+        win = padded.unfold(dimension=2, size=2 * radius + 1, step=1)
+        win = win.permute(0, 2, 3, 1).contiguous()
+        return win
+
+    def _build_context(self, h: torch.Tensor, z: torch.Tensor, b_soft: torch.Tensor) -> torch.Tensor:
+        local = self._build_local_windows(h, self.window_radius)
+        z_exp = z.unsqueeze(2)  # (B,T,1,D)
+        c = torch.cat([z_exp, local], dim=2)  # (B,T,Lc,D)
+        return c
+
+    def _context_logits(self, c: torch.Tensor, comp: str) -> torch.Tensor:
+        # c: (B,T,L,D)
+        scores = self.attn_mlp[comp](c)  # (B,T,L,1)
+        attn = torch.softmax(scores, dim=2)
+        pooled = (attn * c).sum(dim=2)  # (B,T,D)
+        logits = self.comp_proj[comp](pooled)  # (B,T,V)
+        return logits
+
+    def forward_train(
+        self,
+        encoder_input: torch.Tensor,
+        targets: Dict[str, torch.Tensor],
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        target_mask: Optional[torch.Tensor] = None,
+        vocabs: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        device = encoder_input.device
+        h, boundary_logits = self.encoder.encode(encoder_input, src_key_padding_mask)
+        b_soft = self._smooth_boundary(boundary_logits)
+        z = self._apply_film(h, b_soft)
+        c = self._build_context(h, z, b_soft)
+
+        logits = {comp: self._context_logits(c, comp) for comp in self.comp_names}
+        bsz, t_len = targets["quality"].shape
+        if target_mask is None:
+            target_mask = torch.ones(bsz, t_len, dtype=torch.bool, device=device)
+
+        def ce_masked(logits: torch.Tensor, target: torch.Tensor, mask: torch.Tensor, comp_name: str) -> torch.Tensor:
+            num_classes = logits.size(-1)
+            comp_pad = vocabs.get(f"{comp_name}_pad_idx", vocabs["pad_idx"]) if vocabs is not None else 0
+            valid = mask & (target != comp_pad)
+            safe_target = torch.where(valid, target.clamp(min=0, max=num_classes - 1), torch.zeros_like(target))
+            ce = F.cross_entropy(logits.transpose(1, 2), safe_target, reduction="none")
+            denom = valid.float().sum().clamp(min=1.0)
+            return (ce * valid.float()).sum() / denom
+
+        loss_q = ce_masked(logits["quality"], targets["quality"], target_mask, "quality")
+        loss_r = ce_masked(logits["root"], targets["root"], target_mask, "root")
+        loss_b = ce_masked(logits["bass"], targets["bass"], target_mask, "bass")
+        loss_k = torch.tensor(0.0, device=device)
+        if 'key' in self.comp_names and ('key' in targets):
+            loss_k = ce_masked(logits['key'], targets['key'], target_mask, 'key')
+        bce = F.binary_cross_entropy_with_logits(boundary_logits, targets["boundary"].to(boundary_logits.dtype), reduction="none")
+        loss_boundary = (bce * target_mask.float()).sum() / target_mask.float().sum().clamp(min=1.0)
+        total_loss = loss_q + loss_r + loss_b + loss_k + loss_boundary * 3.0
+
+        return {
+            "loss": total_loss,
+            "loss_map": {
+                "quality": loss_q,
+                "root": loss_r,
+                "bass": loss_b,
+                "key": loss_k,
+                "boundary": loss_boundary,
+            },
+            "logits": logits,
+            "boundary_logits": boundary_logits,
+        }
+
+    def forward_infer(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        h, boundary_logits = self.encoder.encode(encoder_input, src_key_padding_mask)
+        b_soft = self._smooth_boundary(boundary_logits)
+        z = self._apply_film(h, b_soft)
+        c = self._build_context(h, z, b_soft)
+        logits = {comp: self._context_logits(c, comp) for comp in self.comp_names}
+        ids = {k: logits[k].argmax(dim=-1) for k in self.comp_names}
+        conf = {f"conf_{k}": logits[k].softmax(dim=-1).max(dim=-1).values for k in self.comp_names}
+        ids.update(conf)
+        ids["boundary"] = boundary_logits
+        return ids
+
+
+class ChordRecognitionModelWrapper(nn.Module):
+    """A thin wrapper around scripts.model.ChordRecognitionModel to normalize inputs."""
+
+    def __init__(self, model_config: Dict[str, Any], vocab_sizes: Dict[str, int]):
+        super().__init__()
+        from .model import ChordRecognitionModel as CR
+
+        self.inner = CR(model_config, vocab_sizes)
+
+    def forward_train(
+        self,
+        encoder_input: torch.Tensor,
+        targets: Dict[str, torch.Tensor],
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        target_mask: Optional[torch.Tensor] = None,
+        vocabs: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        x = _to_patch_input(encoder_input)
+        return self.inner.forward_train(x, targets, src_key_padding_mask, target_mask)
+
+    def forward_infer(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        x = _to_patch_input(encoder_input)
+        return self.inner.forward_infer(x, src_key_padding_mask)
+
+
+class FiLMKTokenKeyModel(nn.Module):
+    """FiLM + KToken decoder with key-conditioned FiLM injection (key used as auxiliary condition only)."""
+
+    def __init__(self, model_config: Dict[str, Any], vocab_sizes: Dict[str, int], key_vocab_size: int):
+        super().__init__()
+        self.config = model_config
+        self.vocab_sizes = vocab_sizes
+        d_model = model_config["d_model"]
+        self.encoder = BaseEncoder(model_config)
+
+        # FiLM
+        self.d_b = max(1, d_model // 4)
+        k_b = int(model_config["boundary_kernel"])
+        self.boundary_smoother = nn.Conv1d(1, 1, kernel_size=k_b, padding=k_b // 2, bias=True)
+        self.boundary_e0 = nn.Parameter(torch.zeros(self.d_b))
+        self.boundary_e1 = nn.Parameter(torch.randn(self.d_b) * 0.02)
+
+        # Key soft embedding
+        self.d_k = max(1, d_model // 8)
+        self.key_embed = nn.Embedding(key_vocab_size, self.d_k)
+        self.key_head = nn.Linear(d_model, key_vocab_size)
+
+        # FiLM projection with concatenated boundary/key embeddings
+        self.film_ln_in = nn.LayerNorm(d_model + self.d_b + self.d_k)
+        self.film_ln_h = nn.LayerNorm(d_model)
+        self.film_mlp = nn.Linear(d_model + self.d_b + self.d_k, 2 * d_model)
+
+        # Decoder (same as CR model)
+        self.mask_id_q = self.vocab_sizes["quality"]
+        self.mask_id_r = self.vocab_sizes["root"]
+        self.mask_id_b = self.vocab_sizes["bass"]
+        self.emb_q = nn.Embedding(self.vocab_sizes["quality"] + 1, d_model)
+        self.emb_r = nn.Embedding(self.vocab_sizes["root"] + 1, d_model)
+        self.emb_b = nn.Embedding(self.vocab_sizes["bass"] + 1, d_model)
+
+        dec_heads = int(model_config["dec_heads"])
+        dec_mlp_ratio = int(model_config["dec_mlp_ratio"])
+        dec_layers = int(model_config["dec_layers"])
+        dec_dropout = float(model_config["dec_dropout"])
+        self.window_radius = int(model_config["window_radius"])
+        self.decoder_layers = nn.ModuleList(
+            [
+                KTokenDecoderLayer(
+                    d_model=d_model,
+                    nhead=dec_heads,
+                    mlp_ratio=dec_mlp_ratio,
+                    dropout=dec_dropout,
+                )
+                for _ in range(dec_layers)
+            ]
+        )
+        self.dec_norm = nn.LayerNorm(d_model)
+        self.head_q = nn.Linear(d_model, self.vocab_sizes["quality"])
+        self.head_r = nn.Linear(d_model, self.vocab_sizes["root"])
+        self.head_b = nn.Linear(d_model, self.vocab_sizes["bass"])
+        self.use_key = ("key" in self.vocab_sizes)
+
+    def _smooth_boundary(self, boundary_logits: torch.Tensor) -> torch.Tensor:
+        b = boundary_logits.unsqueeze(1)
+        smoothed = self.boundary_smoother(b)
+        return torch.sigmoid(smoothed.squeeze(1))
+
+    def _apply_film(self, h: torch.Tensor, b_soft: torch.Tensor, key_soft: torch.Tensor) -> torch.Tensor:
+        bsz, t_len, d_model = h.shape
+        # boundary embedding
+        e0 = self.boundary_e0.to(h.dtype).view(1, 1, -1).expand(bsz, t_len, -1)
+        e1 = self.boundary_e1.to(h.dtype).view(1, 1, -1).expand(bsz, t_len, -1)
+        eb = b_soft.unsqueeze(-1) * e1 + (1.0 - b_soft).unsqueeze(-1) * e0
+        # key soft embedding via expectation over embeddings
+        key_indices = torch.arange(self.key_embed.num_embeddings, device=h.device)
+        key_emb_table = self.key_embed(key_indices)  # (V_k, d_k)
+        key_emb_table = key_emb_table.to(dtype=key_soft.dtype)
+        ek = key_soft @ key_emb_table  # (B,T,d_k)
+
+        film_in = torch.cat([h, eb, ek], dim=-1)
+        film_in = self.film_ln_in(film_in)
+        gamma, beta = self.film_mlp(film_in).chunk(2, dim=-1)
+        z = self.film_ln_h(h) * (1.0 + gamma) + beta
+        return z
+
+    def _build_local_windows(self, h: torch.Tensor, radius: int) -> torch.Tensor:
+        x = h.transpose(1, 2)
+        padded = F.pad(x, (radius, radius), mode="replicate")
+        win = padded.unfold(dimension=2, size=2 * radius + 1, step=1)
+        win = win.permute(0, 2, 3, 1).contiguous()
+        return win
+
+    def _build_context(self, h: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        local = self._build_local_windows(h, self.window_radius)
+        z_exp = z.unsqueeze(2)
+        c = torch.cat([z_exp, local], dim=2)
+        return c
+
+    def _embed_tokens(self, ids_q: torch.Tensor, ids_r: torch.Tensor, ids_b: torch.Tensor) -> torch.Tensor:
+        xq = self.emb_q(ids_q)
+        xr = self.emb_r(ids_r)
+        xb = self.emb_b(ids_b)
+        x = torch.stack([xq, xr, xb], dim=2)
+        return x
+
+    def _run_decoder(self, X: torch.Tensor, C: torch.Tensor):
+        x = X
+        for layer in self.decoder_layers:
+            x = layer(x, C)
+        x = self.dec_norm(x)
+        xq = x[:, :, 0, :]
+        xr = x[:, :, 1, :]
+        xb = x[:, :, 2, :]
+        logits_q = self.head_q(xq)
+        logits_r = self.head_r(xr)
+        logits_b = self.head_b(xb)
+        return logits_q, logits_r, logits_b
+
+    def forward_train(
+        self,
+        encoder_input: torch.Tensor,
+        targets: Dict[str, torch.Tensor],
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        target_mask: Optional[torch.Tensor] = None,
+        vocabs: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        device = encoder_input.device
+        h, boundary_logits = self.encoder.encode(encoder_input, src_key_padding_mask)
+        b_soft = self._smooth_boundary(boundary_logits)
+
+        key_logits = self.key_head(h)
+        key_soft = key_logits.softmax(dim=-1)
+
+        z = self._apply_film(h, b_soft, key_soft)
+        C = self._build_context(h, z)
+
+        tgt_q = targets["quality"]
+        tgt_r = targets["root"]
+        tgt_b = targets["bass"]
+        bsz, t_len = tgt_q.shape
+        if target_mask is None:
+            target_mask = torch.ones(bsz, t_len, dtype=torch.bool, device=device)
+
+        # Random mask-filling per CR model
+        k_rand = torch.randint(1, 4, (bsz, t_len), device=device)
+        rand_scores = torch.rand(bsz, t_len, 3, device=device)
+        top_vals, top_idx = torch.topk(rand_scores, k=3, dim=-1)
+        mask_slots = torch.zeros(bsz, t_len, 3, dtype=torch.bool, device=device)
+        for kk in range(1, 4):
+            sel = (k_rand == kk)
+            if sel.any():
+                idx_sel = top_idx[sel][:, :kk]
+                row = mask_slots[sel]
+                if idx_sel.numel() > 0:
+                    row.scatter_(dim=1, index=idx_sel, value=True)
+                mask_slots[sel] = row
+
+        ids_q = tgt_q.clone()
+        ids_r = tgt_r.clone()
+        ids_b = tgt_b.clone()
+        ids_q[mask_slots[:, :, 0]] = self.mask_id_q
+        ids_r[mask_slots[:, :, 1]] = self.mask_id_r
+        ids_b[mask_slots[:, :, 2]] = self.mask_id_b
+
+        X = self._embed_tokens(ids_q, ids_r, ids_b)
+        logits_q, logits_r, logits_b = self._run_decoder(X, C)
+
+        def ce_masked(logits: torch.Tensor, target: torch.Tensor, slot_mask: torch.Tensor) -> torch.Tensor:
+            m = slot_mask & target_mask
+            num_classes = logits.size(-1)
+            safe_target = torch.where(m, target.clamp(min=0, max=num_classes - 1), torch.zeros_like(target))
+            ce = F.cross_entropy(logits.transpose(1, 2), safe_target, reduction="none")
+            denom = m.float().sum().clamp(min=1.0)
+            return (ce * m.float()).sum() / denom
+
+        loss_q = ce_masked(logits_q, tgt_q, mask_slots[:, :, 0])
+        loss_r = ce_masked(logits_r, tgt_r, mask_slots[:, :, 1])
+        loss_b = ce_masked(logits_b, tgt_b, mask_slots[:, :, 2])
+
+        # Optional key cross-entropy (auxiliary, weight 1.0)
+        loss_k = torch.tensor(0.0, device=device)
+        if self.use_key and ("key" in targets):
+            key_logits = self.key_head(h)
+            num_classes_k = key_logits.size(-1)
+            safe_k = targets["key"].clamp(min=0, max=num_classes_k - 1)
+            loss_k = F.cross_entropy(key_logits.transpose(1, 2), safe_k, reduction="none")
+            denom_k = target_mask.float().sum().clamp(min=1.0)
+            loss_k = (loss_k * target_mask.float()).sum() / denom_k
+
+        bce = F.binary_cross_entropy_with_logits(boundary_logits, targets["boundary"].to(boundary_logits.dtype), reduction="none")
+        loss_boundary = (bce * target_mask.float()).sum() / target_mask.float().sum().clamp(min=1.0)
+        total_loss = loss_q + loss_r + loss_b + loss_k + loss_boundary * 3.0
+
+        return {
+            "loss": total_loss,
+            "loss_map": {
+                "quality": loss_q,
+                "root": loss_r,
+                "bass": loss_b,
+                "key": loss_k,
+                "boundary": loss_boundary,
+            },
+            "logits": {
+                "quality": logits_q,
+                "root": logits_r,
+                "bass": logits_b,
+            },
+            "mask_slots": mask_slots,
+            "boundary_logits": boundary_logits,
+        }
+
+    def forward_infer(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        device = encoder_input.device
+        h, boundary_logits = self.encoder.encode(encoder_input, src_key_padding_mask)
+        b_soft = self._smooth_boundary(boundary_logits)
+        key_logits = self.key_head(h)
+        key_soft = key_logits.softmax(dim=-1)
+        z = self._apply_film(h, b_soft, key_soft)
+        C = self._build_context(h, z)
+
+        bsz, t_len, _ = h.shape
+        ids_q = torch.full((bsz, t_len), self.mask_id_q, dtype=torch.long, device=device)
+        ids_r = torch.full((bsz, t_len), self.mask_id_r, dtype=torch.long, device=device)
+        ids_b = torch.full((bsz, t_len), self.mask_id_b, dtype=torch.long, device=device)
+        filled_q = torch.zeros((bsz, t_len), dtype=torch.bool, device=device)
+        filled_r = torch.zeros((bsz, t_len), dtype=torch.bool, device=device)
+        filled_b = torch.zeros((bsz, t_len), dtype=torch.bool, device=device)
+
+        for step in (3, 2, 1):
+            X = self._embed_tokens(ids_q, ids_r, ids_b)
+            logits_q, logits_r, logits_b = self._run_decoder(X, C)
+            pq = logits_q.softmax(dim=-1)
+            pr = logits_r.softmax(dim=-1)
+            pb = logits_b.softmax(dim=-1)
+            conf_q = pq.max(dim=-1).values
+            conf_r = pr.max(dim=-1).values
+            conf_b = pb.max(dim=-1).values
+            conf_q = conf_q.masked_fill(filled_q, float("-inf"))
+            conf_r = conf_r.masked_fill(filled_r, float("-inf"))
+            conf_b = conf_b.masked_fill(filled_b, float("-inf"))
+            conf = torch.stack([conf_q, conf_r, conf_b], dim=-1)
+            take_slot = conf.argmax(dim=-1)
+
+            pred_q = logits_q.argmax(dim=-1)
+            commit_q = (take_slot == 0) | ((step == 1) & (~filled_q))
+            ids_q[commit_q] = pred_q[commit_q]
+            filled_q = filled_q | commit_q
+
+            pred_r = logits_r.argmax(dim=-1)
+            commit_r = (take_slot == 1) | ((step == 1) & (~filled_r))
+            ids_r[commit_r] = pred_r[commit_r]
+            filled_r = filled_r | commit_r
+
+            pred_b = logits_b.argmax(dim=-1)
+            commit_b = (take_slot == 2) | ((step == 1) & (~filled_b))
+            ids_b[commit_b] = pred_b[commit_b]
+            filled_b = filled_b | commit_b
+
+        X = self._embed_tokens(ids_q, ids_r, ids_b)
+        logits_q, logits_r, logits_b = self._run_decoder(X, C)
+        conf_q = logits_q.softmax(dim=-1).max(dim=-1).values
+        conf_r = logits_r.softmax(dim=-1).max(dim=-1).values
+        conf_b = logits_b.softmax(dim=-1).max(dim=-1).values
+        return {
+            "quality": ids_q,
+            "root": ids_r,
+            "bass": ids_b,
+            "conf_quality": conf_q,
+            "conf_root": conf_r,
+            "conf_bass": conf_b,
+            "boundary": boundary_logits,
+        }
+
+
+class HTAdapter(nn.Module):
+    """Adapter to unify HT with the common training/eval interface."""
+
+    def __init__(self, ht_config: Dict[str, Any], vocab_sizes: Dict[str, int]):
+        super().__init__()
+        from .HT import HarmonyTransformer
+
+        self.inner = HarmonyTransformer(ht_config, vocab_sizes)
+        self.vocab_sizes = vocab_sizes
+
+    def forward_train(
+        self,
+        encoder_input: torch.Tensor,
+        targets: Dict[str, torch.Tensor],
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        target_mask: Optional[torch.Tensor] = None,
+        vocabs: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        device = encoder_input.device
+        out = self.inner(encoder_input, src_key_padding_mask)
+        logits_q = out["quality"]
+        logits_r = out["root"]
+        logits_b = out["bass"]
+        boundary_logits = out["boundary"].squeeze(-1)
+
+        bsz, t_len, _ = logits_q.shape
+        if target_mask is None:
+            target_mask = torch.ones(bsz, t_len, dtype=torch.bool, device=device)
+
+        def ce_masked(logits: torch.Tensor, target: torch.Tensor, mask: torch.Tensor, comp_name: str) -> torch.Tensor:
+            num_classes = logits.size(-1)
+            comp_pad = vocabs.get(f"{comp_name}_pad_idx", vocabs["pad_idx"]) if vocabs is not None else 0
+            valid = mask & (target != comp_pad)
+            safe_target = torch.where(valid, target.clamp(min=0, max=num_classes - 1), torch.zeros_like(target))
+            ce = F.cross_entropy(logits.transpose(1, 2), safe_target, reduction="none")
+            denom = valid.float().sum().clamp(min=1.0)
+            return (ce * valid.float()).sum() / denom
+
+        loss_q = ce_masked(logits_q, targets["quality"], target_mask, "quality")
+        loss_r = ce_masked(logits_r, targets["root"], target_mask, "root")
+        loss_b = ce_masked(logits_b, targets["bass"], target_mask, "bass")
+        bce = F.binary_cross_entropy_with_logits(boundary_logits, targets["boundary"].to(boundary_logits.dtype), reduction="none")
+        loss_boundary = (bce * target_mask.float()).sum() / target_mask.float().sum().clamp(min=1.0)
+        total_loss = loss_q + loss_r + loss_b + 3.0 * loss_boundary
+
+        return {
+            "loss": total_loss,
+            "loss_map": {
+                "quality": loss_q,
+                "root": loss_r,
+                "bass": loss_b,
+                "boundary": loss_boundary,
+            },
+            "logits": {
+                "quality": logits_q,
+                "root": logits_r,
+                "bass": logits_b,
+            },
+            "boundary_logits": boundary_logits,
+        }
+
+    def forward_infer(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        out = self.inner(encoder_input, src_key_padding_mask)
+        ids = {
+            "quality": out["quality"].argmax(dim=-1),
+            "root": out["root"].argmax(dim=-1),
+            "bass": out["bass"].argmax(dim=-1),
+            "boundary": out["boundary"].squeeze(-1),
+        }
+        return ids
+
+
+def build_model(experiment: str, model_config: Dict[str, Any], vocabs: Dict[str, Any], use_key: bool = False) -> nn.Module:
+    # Build vocab sizes for components (include key optionally)
+    chord_components = ["root", "quality", "bass"] + (["key"] if use_key and ("key" in vocabs) else [])
+    vocab_sizes = {comp: len(vocabs[comp]) for comp in chord_components}
+
+    exp = experiment.lower()
+    if exp == "baseline":
+        return BaselineLinearModel(model_config, vocab_sizes)
+    if exp == "film_ctx":
+        return FiLMContextLinearModel(model_config, vocab_sizes)
+    if exp == "film_kdec":
+        return ChordRecognitionModelWrapper(model_config, vocab_sizes)
+    if exp == "film_kdec_key":
+        key_vocab_size = len(vocabs["key"]) if ("key" in vocabs) else 24
+        return FiLMKTokenKeyModel(model_config, vocab_sizes, key_vocab_size)
+    if exp == "ht":
+        # The HT config expects specific keys; reuse model_config where possible
+        ht_cfg = {
+            "input_size": model_config["input_size"],
+            "d_model": model_config["d_model"],
+            "n_layers": model_config["num_encoder_layers"],
+            "n_heads": model_config["n_head"],
+            "dropout": model_config["dropout"],
+            "train_boundary": True,
+            "slope": 1.0,
+            "n_beats": model_config["n_beats"],
+            "beat_resolution": model_config["beat_resolution"],
+        }
+        return HTAdapter(ht_cfg, vocab_sizes)
+
+    raise ValueError(f"Unknown experiment '{experiment}'")
+
+

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,10 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0
+gradio==5.28.0
+torch
+torchaudio
+numpy
+pyyaml
+music21
+miditoolkit
+huggingface_hub
+tqdm