import re import hashlib import json import mir_eval.chord import numpy as np from .durations_sheetsage2 import DURATION_TEMPLATES, duration_boundaries from .labels_sheetsage2 import STRUCTURE_LABELS from .schema_sheetsage2 import get_prompt_multitask_schema PROMPT_ORDER = tuple(task.name for task in get_prompt_multitask_schema("v1").tasks) CHROMATIC_SHARPS = ( "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B", ) FULL_CHORD_QUALITIES = ( "maj", "min", "dim", "aug", "maj7", "min7", "7", "hdim7", "dim7", "minmaj7", "sus2", "sus4", "sus4(b7)", "maj6", "min6", ) FULL_CHORD_INVERSIONS = { "maj": ("/2", "/3", "/5"), "min": ("/2", "/b3", "/5"), "maj7": ("/3", "/5", "/7"), "min7": ("/b3", "/5", "/b7"), "7": ("/3", "/5", "/b7"), } def build_full_chord_vocabulary(): labels = ["N"] for quality in FULL_CHORD_QUALITIES: inversions = FULL_CHORD_INVERSIONS.get(quality, ()) for root in CHROMATIC_SHARPS: for inversion in (*inversions, ""): labels.append(f"{root}:{quality}{inversion}") return tuple(labels) FULL_CHORD_VOCABULARY = build_full_chord_vocabulary() class SheetSage2Tokenizer: """Typed prompt and event vocabulary for prompt-conditioned transcription.""" meter_numerators = tuple(range(1, 33)) meter_denominators = (1, 2, 4, 8, 16, 32) n_eighth_positions = 256 max_subbeat_shift = 256 prompt_capacity = 256 def __init__( self, audio_length_seconds=300.0, time_hz=100, schema_version="v1", expected_fingerprint=None, ): self.audio_length_seconds = float(audio_length_seconds) self.time_hz = int(time_hz) self.schema = get_prompt_multitask_schema(schema_version) self.schema_version = self.schema.version self.task_specs = self.schema.tasks self.event_field_order = self.schema.event_field_order self.n_time_tokens = int(round(self.audio_length_seconds * self.time_hz)) if self.n_time_tokens <= 0: raise ValueError("audio_length_seconds must produce at least one time token") self.pad_token = 0 self.sos_token = 1 self.bos_token = self.sos_token self.eos_token = 2 self.out_token = 3 self.prompt_token_start = 4 self.prompt_names = tuple(task.name for task in self.task_specs) if len(self.prompt_names) > self.prompt_capacity: raise ValueError( f"Schema {self.schema_version} has {len(self.prompt_names)} prompts, " f"exceeding immutable capacity {self.prompt_capacity}" ) self.prompt_to_id = { name: self.prompt_token_start + index for index, name in enumerate(self.prompt_names) } self.prompt_token_end = self.prompt_token_start + self.prompt_capacity self.subbeat_shift_token_start = self.prompt_token_end self.n_subbeat_shift_tokens = self.max_subbeat_shift + 1 self.subbeat_shift_token_end = ( self.subbeat_shift_token_start + self.n_subbeat_shift_tokens ) self.time_token_start = self.subbeat_shift_token_end self.time_token_end = self.time_token_start + self.n_time_tokens self.meter_pairs = tuple( (numerator, denominator) for numerator in self.meter_numerators for denominator in self.meter_denominators ) self.meter_to_id = {meter: index for index, meter in enumerate(self.meter_pairs)} self.meter_token_start = self.time_token_end self.meter_token_end = self.meter_token_start + len(self.meter_pairs) self.eighth_position_token_start = self.meter_token_end self.eighth_position_token_end = ( self.eighth_position_token_start + self.n_eighth_positions ) self.structure_labels = tuple(STRUCTURE_LABELS) self.structure_token_start = self.eighth_position_token_end self.structure_token_end = self.structure_token_start + len(self.structure_labels) self.key_token_start = self.structure_token_end self.n_key_tokens = 24 self.key_token_end = self.key_token_start + self.n_key_tokens self.majmin_chord_labels = ( "N", *(f"{root}:maj" for root in CHROMATIC_SHARPS), *(f"{root}:min" for root in CHROMATIC_SHARPS), ) self.majmin_chord_token_start = self.key_token_end self.majmin_chord_token_end = ( self.majmin_chord_token_start + len(self.majmin_chord_labels) ) self.full_chord_labels = FULL_CHORD_VOCABULARY self.full_chord_to_id = { label: index for index, label in enumerate(self.full_chord_labels) } self.full_chord_token_start = self.majmin_chord_token_end self.full_chord_token_end = ( self.full_chord_token_start + len(self.full_chord_labels) ) self.pitch_token_start = self.full_chord_token_end self.n_pitch_tokens = 256 self.pitch_token_end = self.pitch_token_start + self.n_pitch_tokens self.duration_templates = DURATION_TEMPLATES self.duration_boundaries = duration_boundaries self.duration_token_start = self.pitch_token_end self.n_duration_tokens = len(self.duration_templates) self.duration_token_end = self.duration_token_start + self.n_duration_tokens self.appended_token_blocks = {} next_token = self.duration_token_end for block in self.schema.appended_token_blocks: start = next_token end = start + len(block.labels) self.appended_token_blocks[block.name] = { "start": start, "end": end, "labels": tuple(block.labels), "output_field": block.output_field or block.name, } next_token = end self.n_tokens = next_token self._full_chord_templates = self._build_full_chord_templates() self.vocab_fingerprint = self._compute_fingerprint() if ( expected_fingerprint is not None and str(expected_fingerprint) != self.vocab_fingerprint ): raise ValueError( f"Tokenizer fingerprint mismatch for schema {self.schema_version}: " f"expected {expected_fingerprint}, got {self.vocab_fingerprint}" ) def _compute_fingerprint(self): payload = { "schema_version": self.schema_version, "audio_length_seconds": self.audio_length_seconds, "time_hz": self.time_hz, "prompt_capacity": self.prompt_capacity, "prompt_names": self.prompt_names, "event_field_order": self.event_field_order, "meter_pairs": self.meter_pairs, "structure_labels": self.structure_labels, "majmin_chord_labels": self.majmin_chord_labels, "full_chord_labels": self.full_chord_labels, "duration_templates": tuple(int(value) for value in self.duration_templates), "appended_token_blocks": tuple( (name, block["labels"], block["output_field"]) for name, block in self.appended_token_blocks.items() ), "n_tokens": self.n_tokens, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode( "utf-8" ) return hashlib.sha256(encoded).hexdigest()[:16] def get_config(self): return { "schema_version": self.schema_version, "audio_length_seconds": self.audio_length_seconds, "time_hz": self.time_hz, "vocab_fingerprint": self.vocab_fingerprint, "n_tokens": self.n_tokens, } def normalize_prompts(self, prompts): names = [] seen = set() for prompt in prompts: name = str(prompt).strip() if name.startswith("<|") and name.endswith("|>"): name = name[2:-2] if name not in self.prompt_to_id: raise ValueError(f"Unknown prompt: {prompt!r}") if name not in seen: names.append(name) seen.add(name) names.sort(key=self.prompt_names.index) selected_groups = {} task_by_name = {task.name: task for task in self.task_specs} for name in names: group = task_by_name[name].sampling_group previous = selected_groups.get(group) if previous is not None: raise ValueError( f"Prompts {previous!r} and {name!r} are mutually exclusive " f"within sampling group {group!r}" ) selected_groups[group] = name if not names: raise ValueError("At least one task prompt is required") return tuple(names) def prompt_to_token(self, prompt): return self.prompt_to_id[self.normalize_prompts((prompt,))[0]] def token_to_prompt(self, token): token = int(token) for prompt, prompt_token in self.prompt_to_id.items(): if token == prompt_token: return prompt raise ValueError(f"token {token} is not a prompt token") def prompt_prefix(self, prompts): prompts = self.normalize_prompts(prompts) return [ self.sos_token, *(self.prompt_to_id[prompt] for prompt in prompts), self.out_token, ] def subbeat_shift_to_tokens(self, shift): shift = int(shift) if shift < 0: raise ValueError("subbeat shift must be non-negative") tokens = [] while shift > self.max_subbeat_shift: tokens.append(self.subbeat_shift_token_start + self.max_subbeat_shift) shift -= self.max_subbeat_shift tokens.append(self.subbeat_shift_token_start + shift) return tokens def token_to_subbeat_shift(self, token): token = int(token) if not self.subbeat_shift_token_start <= token < self.subbeat_shift_token_end: raise ValueError(f"token {token} is not a subbeat shift token") return token - self.subbeat_shift_token_start def time_id_to_token(self, time_id): time_id = int(time_id) if not 0 <= time_id < self.n_time_tokens: raise ValueError(f"time id {time_id} is outside [0, {self.n_time_tokens})") return self.time_token_start + time_id def token_to_time_id(self, token): token = int(token) if not self.time_token_start <= token < self.time_token_end: raise ValueError(f"token {token} is not a time token") return token - self.time_token_start def meter_to_token(self, numerator, denominator): meter = (int(numerator), int(denominator)) if meter not in self.meter_to_id: raise ValueError(f"Unsupported meter {meter[0]}/{meter[1]}") return self.meter_token_start + self.meter_to_id[meter] def eighth_position_to_token(self, position): position = int(position) if not 0 <= position < self.n_eighth_positions: raise ValueError( f"eighth-note position {position} is outside [0, {self.n_eighth_positions})" ) return self.eighth_position_token_start + position def structure_to_token(self, label): label = str(label) if label not in self.structure_labels: raise ValueError(f"Unknown structure label: {label!r}") return self.structure_token_start + self.structure_labels.index(label) def key_to_token(self, key_label, pitch_shift=0): tonic, mode = str(key_label).split(":", 1) tonic_id = mir_eval.chord.pitch_class_to_semitone(tonic) if tonic_id < 0: raise ValueError(f"Invalid key tonic: {key_label!r}") tonic_id = (int(tonic_id) + int(pitch_shift)) % 12 minor_modes = {"minor", "dorian", "phrygian", "locrian"} mode_id = 1 if mode.lower() in minor_modes else 0 return self.key_token_start + mode_id * 12 + tonic_id @staticmethod def _canonical_chord_parts(chord_label, pitch_shift=0): chord_label = str(chord_label).strip() if chord_label in {"N", "X", ""}: return "N", None, None root, suffix = chord_label.split(":", 1) root_id = mir_eval.chord.pitch_class_to_semitone(root) if root_id < 0: return "N", None, None root_id = (int(root_id) + int(pitch_shift)) % 12 canonical = f"{CHROMATIC_SHARPS[root_id]}:{suffix}" quality = re.split(r"/", suffix, maxsplit=1)[0] return canonical, root_id, quality def chord_majmin_to_token(self, chord_label, pitch_shift=0): canonical, root_id, quality = self._canonical_chord_parts( chord_label, pitch_shift, ) if canonical == "N": return self.majmin_chord_token_start try: original_root, chroma, _ = mir_eval.chord.encode(str(chord_label)) relative = mir_eval.chord.rotate_bitmap_to_root(chroma, original_root) has_minor_third = bool(relative[3] > 0) has_major_third = bool(relative[4] > 0) except Exception: has_minor_third = quality.startswith(("min", "dim", "hdim")) has_major_third = not has_minor_third if has_major_third and not has_minor_third: quality_id = 0 elif has_minor_third and not has_major_third: quality_id = 1 elif quality.startswith(("min", "dim", "hdim")): quality_id = 1 elif has_major_third: quality_id = 0 else: return self.majmin_chord_token_start return self.majmin_chord_token_start + 1 + quality_id * 12 + root_id @staticmethod def _chord_template(chord_label): root, chroma, bass = mir_eval.chord.encode(chord_label) root_chroma = np.zeros(12, dtype=np.float32) root_chroma[root] = 1.0 relative_chroma = mir_eval.chord.rotate_bitmap_to_root(chroma, root).astype( np.float32 ) bass_chroma = np.zeros(12, dtype=np.float32) bass_chroma[(bass + root) % 12] = 1.0 return np.concatenate([root_chroma, relative_chroma, bass_chroma]) def _build_full_chord_templates(self): templates = [np.zeros(36, dtype=np.float32)] templates.extend( self._chord_template(label) for label in self.full_chord_labels[1:] ) return np.stack(templates) def chord_full_to_token(self, chord_label, pitch_shift=0): canonical, _root_id, _quality = self._canonical_chord_parts( chord_label, pitch_shift, ) chord_id = self.full_chord_to_id.get(canonical) if chord_id is None: if canonical == "N": chord_id = 0 else: try: target = self._chord_template(canonical) distances = np.abs(self._full_chord_templates - target[None]).sum( axis=1 ) chord_id = int(np.argmin(distances)) except Exception: chord_id = 0 return self.full_chord_token_start + chord_id def pitch_to_token(self, pitch, track=0, full_melody=False): pitch = int(pitch) track = int(track) if not 0 <= pitch < 128: raise ValueError(f"MIDI pitch {pitch} is outside [0, 128)") if track not in (0, 1): raise ValueError(f"melody track {track} is outside [0, 2)") pitch_id = pitch + (128 if full_melody and track == 1 else 0) return self.pitch_token_start + pitch_id def duration_bin_to_token(self, duration_bin): duration_bin = int(duration_bin) if not 0 <= duration_bin < self.n_duration_tokens: raise ValueError( f"duration bin {duration_bin} is outside [0, {self.n_duration_tokens})" ) return self.duration_token_start + duration_bin def token_to_duration_bin(self, token): token = int(token) if not self.duration_token_start <= token < self.duration_token_end: raise ValueError(f"token {token} is not a duration token") return token - self.duration_token_start def appended_label_to_token(self, block_name, label): block = self.appended_token_blocks.get(str(block_name)) if block is None: raise ValueError(f"unknown appended token block: {block_name!r}") try: index = block["labels"].index(str(label)) except ValueError as exc: raise ValueError( f"unknown label {label!r} for appended block {block_name!r}" ) from exc return block["start"] + index def token_to_appended_label(self, token): token = int(token) token_type = self.token_type(token) block = self.appended_token_blocks.get(token_type) if block is None: raise ValueError(f"token {token} is not from an appended token block") return token_type, block["labels"][token - block["start"]] def _token_output_field(self, token_type): built_in = { "time": "timestamp", "meter": "rhythm", "eighth_position": "rhythm", "structure": "structure", "key": "key", "chord_majmin": "chord", "chord_full": "chord", "pitch": "melody", "duration": "melody", } if token_type in built_in: return built_in[token_type] block = self.appended_token_blocks.get(token_type) return None if block is None else block["output_field"] def _decode_field(self, field, field_tokens, prompts): token_types = [self.token_type(token) for token in field_tokens] if field == "timestamp": return self.token_to_time_id(field_tokens[0]) / self.time_hz if field == "rhythm": rhythm = {} for token, token_type in zip(field_tokens, token_types): if token_type == "meter": rhythm["meter"] = self.meter_pairs[token - self.meter_token_start] elif token_type == "eighth_position": rhythm["eighth_position"] = token - self.eighth_position_token_start return rhythm if field == "structure": return self.structure_labels[field_tokens[0] - self.structure_token_start] if field == "key": key_id = field_tokens[0] - self.key_token_start mode = "minor" if key_id >= 12 else "major" return f"{CHROMATIC_SHARPS[key_id % 12]}:{mode}" if field == "chord": if token_types[0] == "chord_majmin": return self.majmin_chord_labels[ field_tokens[0] - self.majmin_chord_token_start ] return self.full_chord_labels[ field_tokens[0] - self.full_chord_token_start ] if field == "melody": notes = [] index = 0 while index < len(field_tokens): pitch_id = field_tokens[index] - self.pitch_token_start duration_bin = 0 if ( index + 1 < len(field_tokens) and self.token_type(field_tokens[index + 1]) == "duration" ): duration_bin = self.token_to_duration_bin(field_tokens[index + 1]) index += 2 else: index += 1 notes.append( { "pitch": pitch_id % 128, "track": int(pitch_id >= 128), "duration_bin": duration_bin, "duration_steps": int(self.duration_templates[duration_bin]), } ) return notes values = [] for token in field_tokens: block_name, label = self.token_to_appended_label(token) values.append({"block": block_name, "label": label}) return values[0] if len(values) == 1 else values def decode_sequence(self, tokens, strict=True): """Parse one prompt-conditioned sequence into timed, typed events.""" if hasattr(tokens, "detach"): tokens = tokens.detach().cpu().tolist() tokens = [int(token) for token in tokens] while tokens and tokens[-1] == self.pad_token: tokens.pop() if not tokens or tokens[0] != self.sos_token: raise ValueError("sequence must begin with <|sos|>") try: out_index = tokens.index(self.out_token, 1) except ValueError as exc: raise ValueError("sequence is missing <|out|>") from exc prompts = tuple(self.token_to_prompt(token) for token in tokens[1:out_index]) if strict and self.normalize_prompts(prompts) != prompts: raise ValueError("prompt tokens are not in canonical schema order") active_fields = { task.output_field for task in self.task_specs if task.name in prompts } events = [] position = out_index + 1 current_step = 0 saw_eos = False while position < len(tokens): token = tokens[position] if token == self.eos_token: saw_eos = True position += 1 break if self.token_type(token) != "subbeat_shift": raise ValueError(f"event at token index {position} has no subbeat shift") shift = 0 while ( position < len(tokens) and self.token_type(tokens[position]) == "subbeat_shift" ): shift += self.token_to_subbeat_shift(tokens[position]) position += 1 current_step += shift tokens_by_field = {field: [] for field in self.event_field_order} while position < len(tokens): token = tokens[position] token_type = self.token_type(token) if token_type == "subbeat_shift" or token == self.eos_token: break field = self._token_output_field(token_type) if field is None or field not in tokens_by_field: raise ValueError( f"token {token} ({token_type}) has no field in schema " f"{self.schema_version}" ) if strict and field not in active_fields: raise ValueError( f"token {token} belongs to inactive output field {field!r}" ) tokens_by_field[field].append(token) position += 1 tokens_by_field = { field: values for field, values in tokens_by_field.items() if values } if not tokens_by_field: if strict: raise ValueError(f"empty event at subbeat {current_step}") continue if strict: for field, values in tokens_by_field.items(): types = [self.token_type(value) for value in values] if field == "timestamp" and types != ["time"]: raise ValueError("timestamp event must contain exactly one time token") if field == "rhythm": if types not in (["eighth_position"], ["meter", "eighth_position"]): raise ValueError(f"invalid rhythm payload: {types}") if field in {"structure", "key", "chord"} and len(values) != 1: raise ValueError(f"field {field!r} must contain exactly one token") if field == "melody": index = 0 while index < len(types): if types[index] != "pitch": raise ValueError( "melody payload must contain pitch tokens with optional duration" ) if index + 1 < len(types) and types[index + 1] == "duration": index += 2 else: index += 1 events.append( { "subbeat": current_step, "tokens_by_field": tokens_by_field, "values": { field: self._decode_field(field, values, prompts) for field, values in tokens_by_field.items() }, } ) if strict and not saw_eos: raise ValueError("sequence is missing <|eos|>") if strict and position != len(tokens): raise ValueError("non-padding tokens follow <|eos|>") return { "schema_version": self.schema_version, "prompts": prompts, "events": events, "has_eos": saw_eos, } def encode_decoded_sequence(self, decoded): """Re-encode the lossless representation returned by decode_sequence.""" prompts = self.normalize_prompts(decoded["prompts"]) output = self.prompt_prefix(prompts) previous_step = 0 for event in decoded["events"]: step = int(event["subbeat"]) if step < previous_step: raise ValueError("events must be sorted by non-decreasing subbeat") output.extend(self.subbeat_shift_to_tokens(step - previous_step)) previous_step = step tokens_by_field = event["tokens_by_field"] for field in self.event_field_order: output.extend(int(token) for token in tokens_by_field.get(field, ())) if decoded.get("has_eos", True): output.append(self.eos_token) return output def token_type(self, token): token = int(token) for name, start, end in ( ("subbeat_shift", self.subbeat_shift_token_start, self.subbeat_shift_token_end), ("time", self.time_token_start, self.time_token_end), ("meter", self.meter_token_start, self.meter_token_end), ("eighth_position", self.eighth_position_token_start, self.eighth_position_token_end), ("structure", self.structure_token_start, self.structure_token_end), ("key", self.key_token_start, self.key_token_end), ("chord_majmin", self.majmin_chord_token_start, self.majmin_chord_token_end), ("chord_full", self.full_chord_token_start, self.full_chord_token_end), ("pitch", self.pitch_token_start, self.pitch_token_end), ("duration", self.duration_token_start, self.duration_token_end), ): if start <= token < end: return name if token in self.prompt_to_id.values(): return "prompt" for name, block in self.appended_token_blocks.items(): if block["start"] <= token < block["end"]: return name if token == self.pad_token: return "pad" if token == self.sos_token: return "sos" if token == self.eos_token: return "eos" if token == self.out_token: return "out" raise ValueError(f"token {token} is outside vocabulary size {self.n_tokens}") def describe(self, token): token = int(token) token_type = self.token_type(token) if token_type == "prompt": prompt_by_id = {value: key for key, value in self.prompt_to_id.items()} return f"<|{prompt_by_id[token]}|>" if token_type == "subbeat_shift": return f"" if token_type == "time": time_id = token - self.time_token_start return f"" if token_type == "meter": numerator, denominator = self.meter_pairs[token - self.meter_token_start] return f"" if token_type == "eighth_position": return f"" if token_type == "structure": return f"" if token_type == "key": key_id = token - self.key_token_start mode = "minor" if key_id >= 12 else "major" return f"" if token_type == "chord_majmin": return f"" if token_type == "chord_full": return f"" if token_type == "pitch": pitch = token - self.pitch_token_start track = 1 if pitch >= 128 else 0 return f"" if token_type == "duration": return f"" if token_type in self.appended_token_blocks: block = self.appended_token_blocks[token_type] label = block["labels"][token - block["start"]] return f"<{token_type}_{label}>" return f"<|{token_type}|>"