| """Prompt grammar, cached generation and overlap context from the evaluated model.""" |
| import copy |
| from contextlib import nullcontext |
| from pathlib import Path |
| import numpy as np |
| import torch |
|
|
| FULL_TASK_PROMPTS = ("timestamp", "downbeat_meter", "structure", "key", "chord_full", "melody_full") |
| FIELD_TO_INDEX = {"timestamp": 0, "rhythm": 1, "structure": 2, "key": 3, "chord": 4, "melody": 5} |
|
|
|
|
| class PromptGrammarState: |
| def __init__(self, tokenizer): |
| self.tokenizer = tokenizer |
| self.generated_events = 0 |
| self.in_shift = True |
| self.shift_run = 0 |
| self.payload_count = 0 |
| self.last_field_index = -1 |
| self.incomplete = None |
|
|
| def _allow_field_starts(self, allowed): |
| tokenizer = self.tokenizer |
| if self.last_field_index < FIELD_TO_INDEX["timestamp"]: |
| allowed[tokenizer.time_token_start : tokenizer.time_token_end] = True |
| if self.last_field_index < FIELD_TO_INDEX["rhythm"]: |
| allowed[tokenizer.meter_token_start : tokenizer.meter_token_end] = True |
| allowed[ |
| tokenizer.eighth_position_token_start : tokenizer.eighth_position_token_end |
| ] = True |
| if self.last_field_index < FIELD_TO_INDEX["structure"]: |
| allowed[ |
| tokenizer.structure_token_start : tokenizer.structure_token_end |
| ] = True |
| if self.last_field_index < FIELD_TO_INDEX["key"]: |
| allowed[tokenizer.key_token_start : tokenizer.key_token_end] = True |
| if self.last_field_index < FIELD_TO_INDEX["chord"]: |
| allowed[ |
| tokenizer.full_chord_token_start : tokenizer.full_chord_token_end |
| ] = True |
| if self.last_field_index <= FIELD_TO_INDEX["melody"]: |
| allowed[tokenizer.pitch_token_start : tokenizer.pitch_token_end] = True |
|
|
| def allowed(self, device): |
| tokenizer = self.tokenizer |
| allowed = torch.zeros(tokenizer.n_tokens, dtype=torch.bool, device=device) |
| can_end = self.payload_count > 0 |
|
|
| if can_end: |
| allowed[tokenizer.eos_token] = True |
| if self.payload_count > 0 or self.in_shift: |
| if self.shift_run < 4: |
| allowed[ |
| tokenizer.subbeat_shift_token_start : tokenizer.subbeat_shift_token_end |
| ] = True |
|
|
| if self.incomplete == "rhythm_after_meter": |
| allowed[ |
| tokenizer.eighth_position_token_start : tokenizer.eighth_position_token_end |
| ] = True |
| return allowed |
|
|
| if self.incomplete == "melody_after_pitch": |
| allowed[tokenizer.duration_token_start : tokenizer.duration_token_end] = True |
| allowed[tokenizer.pitch_token_start : tokenizer.pitch_token_end] = True |
| return allowed |
|
|
| self._allow_field_starts(allowed) |
| return allowed |
|
|
| def update(self, token): |
| tokenizer = self.tokenizer |
| token = int(token) |
| token_type = tokenizer.token_type(token) |
| if token == tokenizer.eos_token: |
| return True |
| if token_type == "subbeat_shift": |
| if not self.in_shift and self.payload_count > 0: |
| self.generated_events += 1 |
| self.payload_count = 0 |
| self.last_field_index = -1 |
| self.incomplete = None |
| self.in_shift = True |
| self.shift_run += 1 |
| return False |
|
|
| self.in_shift = False |
| self.shift_run = 0 |
| self.payload_count += 1 |
| if token_type == "time": |
| self.last_field_index = FIELD_TO_INDEX["timestamp"] |
| self.incomplete = None |
| elif token_type == "meter": |
| self.last_field_index = FIELD_TO_INDEX["rhythm"] |
| self.incomplete = "rhythm_after_meter" |
| elif token_type == "eighth_position": |
| self.last_field_index = FIELD_TO_INDEX["rhythm"] |
| self.incomplete = None |
| elif token_type == "structure": |
| self.last_field_index = FIELD_TO_INDEX["structure"] |
| self.incomplete = None |
| elif token_type == "key": |
| self.last_field_index = FIELD_TO_INDEX["key"] |
| self.incomplete = None |
| elif token_type == "chord_full": |
| self.last_field_index = FIELD_TO_INDEX["chord"] |
| self.incomplete = None |
| elif token_type == "pitch": |
| self.last_field_index = FIELD_TO_INDEX["melody"] |
| self.incomplete = "melody_after_pitch" |
| elif token_type == "duration": |
| self.last_field_index = FIELD_TO_INDEX["melody"] |
| self.incomplete = None |
| else: |
| raise RuntimeError(f"Unexpected prompt token type {token_type!r}") |
| return False |
|
|
|
|
| def inference_autocast(device, dtype=torch.bfloat16): |
| if device.type != "cuda" or dtype is None: |
| return nullcontext() |
| return torch.autocast(device_type="cuda", dtype=dtype) |
|
|
|
|
| def select_cache_batch(cache, indices, batch_size): |
| if cache is None: |
| return None |
| batch_select = getattr(cache, "batch_select_indices", None) |
| if callable(batch_select): |
| batch_select(indices) |
| return cache |
| if torch.is_tensor(cache): |
| if cache.ndim > 0 and cache.shape[0] == batch_size: |
| return cache.index_select(0, indices) |
| return cache |
| if isinstance(cache, tuple): |
| return tuple(select_cache_batch(item, indices, batch_size) for item in cache) |
| if isinstance(cache, list): |
| return [select_cache_batch(item, indices, batch_size) for item in cache] |
| if isinstance(cache, dict): |
| return { |
| key: select_cache_batch(value, indices, batch_size) |
| for key, value in cache.items() |
| } |
| return cache |
|
|
|
|
| @torch.inference_mode() |
| def constrained_prompt_generate_batch( |
| model, |
| audio, |
| prompts, |
| max_sequence_length, |
| prefix_tokens=None, |
| autocast_dtype=torch.bfloat16, |
| stop_time_seconds=None, |
| progress_callback=None, |
| memory=None, |
| step_callback=None, |
| ): |
| tokenizer = model.tokenizer |
| prefix = list(prefix_tokens) if prefix_tokens is not None else tokenizer.prompt_prefix(prompts) |
| if not prefix or prefix[0] != tokenizer.sos_token: |
| raise ValueError("generation prefix must begin with <|sos|>") |
| if prefix[-1] == tokenizer.eos_token: |
| prefix = prefix[:-1] |
| if audio.ndim != 2 or audio.shape[0] < 1: |
| raise ValueError(f"audio must have shape [batch, samples], got {tuple(audio.shape)}") |
| states = [PromptGrammarState(tokenizer) for _ in range(audio.shape[0])] |
| stop_times = [ |
| None if stop_time_seconds is None else float(stop_time_seconds) |
| for _ in states |
| ] |
| out_index = prefix.index(tokenizer.out_token) |
| for state in states: |
| for token in prefix[out_index + 1 :]: |
| state.update(token) |
| output_tokens = [list(prefix) for _ in states] |
| active_sample_ids = list(range(audio.shape[0])) |
| decoder_input = torch.tensor( |
| [prefix] * audio.shape[0], |
| dtype=torch.long, |
| device=audio.device, |
| ) |
| past_key_values = None |
| if memory is None: |
| with inference_autocast(audio.device, autocast_dtype): |
| memory = model.encode(audio) |
| current_length = len(prefix) |
| while active_sample_ids and current_length < int(max_sequence_length): |
| with inference_autocast(audio.device, autocast_dtype): |
| logits, past_key_values = model.decode( |
| memory, |
| decoder_input, |
| use_cache=True, |
| past_key_values=past_key_values, |
| ) |
| next_logits = logits[:, -1].float() |
| allowed = torch.stack( |
| [state.allowed(next_logits.device) for state in states], |
| dim=0, |
| ) |
| masked_logits = next_logits.masked_fill(~allowed, float("-inf")) |
| next_tokens = masked_logits.argmax(dim=-1) |
| if step_callback is not None: |
| step_callback(current_length, active_sample_ids, next_logits, masked_logits) |
| next_token_values = next_tokens.tolist() |
| keep_positions = [] |
| for position, (sample_id, state, stop_time, token) in enumerate( |
| zip(active_sample_ids, states, stop_times, next_token_values) |
| ): |
| output_tokens[sample_id].append(int(token)) |
| finished = state.update(token) |
| if ( |
| not finished |
| and stop_time is not None |
| and tokenizer.time_token_start <= int(token) < tokenizer.time_token_end |
| and tokenizer.token_to_time_id(token) / tokenizer.time_hz >= stop_time |
| ): |
| output_tokens[sample_id].append(tokenizer.eos_token) |
| finished = True |
| if not finished: |
| keep_positions.append(position) |
| current_length += 1 |
| if progress_callback is not None and current_length % 64 == 0: |
| progress_callback(current_length) |
| if not keep_positions: |
| break |
|
|
| old_batch_size = len(active_sample_ids) |
| decoder_input = next_tokens[:, None] |
| if len(keep_positions) == old_batch_size: |
| continue |
|
|
| keep = torch.tensor(keep_positions, dtype=torch.long, device=audio.device) |
| active_sample_ids = [active_sample_ids[position] for position in keep_positions] |
| states = [states[position] for position in keep_positions] |
| stop_times = [stop_times[position] for position in keep_positions] |
| memory = memory.index_select(0, keep) |
| past_key_values = select_cache_batch(past_key_values, keep, old_batch_size) |
| decoder_input = decoder_input.index_select(0, keep) |
|
|
| for tokens in output_tokens: |
| if tokens[-1] != tokenizer.eos_token: |
| tokens.append(tokenizer.eos_token) |
| return [torch.tensor(tokens, dtype=torch.long) for tokens in output_tokens] |
|
|
|
|
| @torch.inference_mode() |
| def constrained_prompt_generate( |
| model, |
| audio, |
| prompts, |
| max_sequence_length, |
| prefix_tokens=None, |
| autocast_dtype=torch.bfloat16, |
| stop_time_seconds=None, |
| progress_callback=None, |
| memory=None, |
| step_callback=None, |
| ): |
| if audio.shape[0] != 1: |
| raise ValueError( |
| "constrained_prompt_generate expects batch size 1; " |
| "use constrained_prompt_generate_batch for batched inference" |
| ) |
| return constrained_prompt_generate_batch( |
| model, |
| audio, |
| prompts, |
| max_sequence_length, |
| prefix_tokens=prefix_tokens, |
| autocast_dtype=autocast_dtype, |
| stop_time_seconds=stop_time_seconds, |
| progress_callback=progress_callback, |
| memory=memory, |
| step_callback=step_callback, |
| )[0] |
|
|
|
|
| def decode_generated_tokens(tokenizer, tokens, song_id, window_index=None): |
| try: |
| return tokenizer.decode_sequence(tokens, strict=True), None |
| except ValueError as exc: |
| recoverable_errors = ( |
| "empty event at subbeat", |
| "belongs to inactive output field", |
| ) |
| if not any(message in str(exc) for message in recoverable_errors): |
| raise |
| decoded = tokenizer.decode_sequence(tokens, strict=False) |
| warning = { |
| "song_id": song_id, |
| "warning": "strict decode failed; recovered in non-strict decode", |
| "error": str(exc), |
| } |
| if window_index is not None: |
| warning["window_index"] = int(window_index) |
| return decoded, warning |
|
|
|
|
| def event_time_map(decoded, target_seconds): |
| anchors = [] |
| for event in decoded["events"]: |
| value = event["values"].get("timestamp") |
| if value is not None: |
| anchors.append((int(event["subbeat"]), float(value))) |
| if not anchors: |
| return lambda step: min(float(target_seconds), max(0.0, float(step) * 0.125)) |
| anchors = sorted(dict(anchors).items()) |
| steps = np.asarray([item[0] for item in anchors], dtype=np.float64) |
| times = np.asarray([item[1] for item in anchors], dtype=np.float64) |
| if len(anchors) >= 2: |
| step_seconds = float(np.median(np.diff(times) / np.maximum(np.diff(steps), 1))) |
| if not np.isfinite(step_seconds) or step_seconds <= 0: |
| step_seconds = 0.125 |
| else: |
| step_seconds = 0.125 |
|
|
| def lookup(step): |
| step = float(step) |
| if step <= steps[0]: |
| return float(np.clip(times[0] + (step - steps[0]) * step_seconds, 0, target_seconds)) |
| if step >= steps[-1]: |
| return float(np.clip(times[-1] + (step - steps[-1]) * step_seconds, 0, target_seconds)) |
| return float(np.interp(step, steps, times)) |
|
|
| return lookup |
|
|
|
|
| def field_text(event, tokenizer): |
| parts = [] |
| for field in tokenizer.event_field_order: |
| value = event["values"].get(field) |
| if value is None: |
| continue |
| if field == "melody": |
| note_parts = [] |
| for note in value: |
| duration = note["duration_bin"] |
| note_parts.append( |
| f"pitch={note['pitch']}:track={note['track']}:dur_bin={duration}:dur_steps={note['duration_steps']}" |
| ) |
| parts.append("melody=[" + ",".join(note_parts) + "]") |
| elif isinstance(value, dict): |
| parts.append(field + "=" + ",".join(f"{k}:{v}" for k, v in value.items())) |
| else: |
| parts.append(f"{field}={value}") |
| return "; ".join(parts) |
|
|
|
|
| def event_start_time(event, time_lookup): |
| if "time" in event: |
| return float(event["time"]) |
| return float(time_lookup(event["subbeat"])) |
|
|
|
|
| def write_events_tsv(decoded, tokenizer, time_lookup, path): |
| with Path(path).open("w", encoding="utf-8") as f: |
| f.write("event_index\tsubbeat\ttime\tfields\ttokens\n") |
| for index, event in enumerate(decoded["events"]): |
| token_text = " ".join( |
| tokenizer.describe(token) |
| for field in tokenizer.event_field_order |
| for token in event["tokens_by_field"].get(field, ()) |
| ) |
| f.write( |
| f"{index}\t{event['subbeat']}\t{event_start_time(event, time_lookup):.6f}\t" |
| f"{field_text(event, tokenizer)}\t{token_text}\n" |
| ) |
|
|
|
|
| def clone_decoded_event(event): |
| return { |
| "subbeat": int(event["subbeat"]), |
| "tokens_by_field": { |
| field: [int(token) for token in tokens] |
| for field, tokens in event["tokens_by_field"].items() |
| }, |
| "values": copy.deepcopy(event["values"]), |
| } |
|
|
|
|
| def refresh_event_values(event, tokenizer, prompts): |
| event["values"] = { |
| field: tokenizer._decode_field(field, tokens, prompts) |
| for field, tokens in event["tokens_by_field"].items() |
| if tokens |
| } |
|
|
|
|
| def set_event_local_timestamp(event, tokenizer, local_time): |
| time_id = int(round(float(local_time) * tokenizer.time_hz)) |
| time_id = max(0, min(time_id, tokenizer.n_time_tokens - 1)) |
| event["tokens_by_field"]["timestamp"] = [tokenizer.time_id_to_token(time_id)] |
|
|
|
|
| def active_context_before(events, tokenizer, time_abs): |
| state = {} |
| for event in events: |
| event_time = event.get("time") |
| if event_time is None or float(event_time) > float(time_abs) + 1e-6: |
| continue |
| for field in ("structure", "key", "chord"): |
| tokens = event["tokens_by_field"].get(field) |
| if tokens: |
| state[field] = [int(token) for token in tokens] |
| rhythm_tokens = event["tokens_by_field"].get("rhythm", ()) |
| meter_tokens = [ |
| int(token) |
| for token in rhythm_tokens |
| if tokenizer.token_type(token) == "meter" |
| ] |
| if meter_tokens: |
| state["meter"] = meter_tokens[:1] |
| return state |
|
|
|
|
| def apply_prefix_context(event, context, tokenizer, prompts): |
| for field in ("structure", "key", "chord"): |
| if field not in event["tokens_by_field"] and field in context: |
| event["tokens_by_field"][field] = list(context[field]) |
|
|
| rhythm_tokens = list(event["tokens_by_field"].get("rhythm", ())) |
| has_meter = any(tokenizer.token_type(token) == "meter" for token in rhythm_tokens) |
| has_eighth = any( |
| tokenizer.token_type(token) == "eighth_position" for token in rhythm_tokens |
| ) |
| if has_eighth and not has_meter and "meter" in context: |
| event["tokens_by_field"]["rhythm"] = list(context["meter"]) + rhythm_tokens |
|
|
| refresh_event_values(event, tokenizer, prompts) |
|
|
|
|
| def build_overlap_prefix_tokens( |
| stitched_events, |
| tokenizer, |
| prompts, |
| window_start, |
| prefix_end, |
| ): |
| eps = 1e-4 |
| source_events = [ |
| event |
| for event in stitched_events |
| if float(window_start) - eps <= float(event.get("time", -1.0)) < float(prefix_end) - eps |
| ] |
| source_events.sort( |
| key=lambda event: ( |
| int( |
| event.get( |
| "global_subbeat", |
| event.get("source_subbeat", event["subbeat"]), |
| ) |
| ), |
| float(event.get("time", 0.0)), |
| ) |
| ) |
| first_beat_index = next( |
| ( |
| index |
| for index, event in enumerate(source_events) |
| if "timestamp" in event["values"] or "rhythm" in event["values"] |
| ), |
| None, |
| ) |
| if first_beat_index is None: |
| return None, None, None |
|
|
| source_events = source_events[first_beat_index:] |
| base_subbeat = int( |
| source_events[0].get( |
| "global_subbeat", |
| source_events[0].get("source_subbeat", source_events[0]["subbeat"]), |
| ) |
| ) |
| context = active_context_before( |
| stitched_events, |
| tokenizer, |
| source_events[0]["time"], |
| ) |
| prefix_events = [] |
| for source in source_events: |
| event = clone_decoded_event(source) |
| source_subbeat = int( |
| source.get( |
| "global_subbeat", |
| source.get("source_subbeat", source["subbeat"]), |
| ) |
| ) |
| event["subbeat"] = max(0, source_subbeat - base_subbeat) |
| if "timestamp" in event["tokens_by_field"]: |
| set_event_local_timestamp( |
| event, |
| tokenizer, |
| float(source["time"]) - float(window_start), |
| ) |
| refresh_event_values(event, tokenizer, prompts) |
| prefix_events.append(event) |
|
|
| apply_prefix_context(prefix_events[0], context, tokenizer, prompts) |
| prefix_decoded = { |
| "schema_version": tokenizer.schema_version, |
| "prompts": prompts, |
| "events": prefix_events, |
| "has_eos": False, |
| } |
| return ( |
| prefix_decoded, |
| tokenizer.encode_decoded_sequence(prefix_decoded), |
| base_subbeat, |
| ) |
|
|
|
|
| def stitched_window_events( |
| decoded, |
| time_lookup, |
| window_start, |
| accept_start, |
| accept_end, |
| song_duration, |
| window_index, |
| global_subbeat_base=0, |
| ): |
| accepted = [] |
| eps = 1e-4 |
| for event in decoded["events"]: |
| local_time = float(time_lookup(event["subbeat"])) |
| abs_time = float(window_start) + local_time |
| if abs_time < float(accept_start) - eps: |
| continue |
| if abs_time >= float(accept_end) - eps or abs_time >= float(song_duration) - eps: |
| continue |
|
|
| output = clone_decoded_event(event) |
| output["time"] = float(np.clip(abs_time, 0.0, song_duration)) |
| output["window_index"] = int(window_index) |
| output["window_start"] = float(window_start) |
| output["source_subbeat"] = int(event["subbeat"]) |
| output["global_subbeat"] = int(global_subbeat_base) + int(event["subbeat"]) |
| if "timestamp" in output["values"]: |
| output["values"]["timestamp"] = output["time"] |
|
|
| notes = output["values"].get("melody") |
| if notes is not None: |
| fixed_notes = [] |
| for note in notes: |
| fixed_note = copy.deepcopy(note) |
| duration_steps = int(fixed_note["duration_steps"]) |
| local_end = float(time_lookup(int(event["subbeat"]) + duration_steps)) |
| end_time = float(window_start) + local_end |
| end_time = min( |
| float(song_duration), |
| max(output["time"] + 0.04, end_time), |
| ) |
| fixed_note["end_time"] = end_time |
| fixed_notes.append(fixed_note) |
| output["values"]["melody"] = fixed_notes |
| accepted.append(output) |
| return accepted |
|
|
|
|
| def write_window_tokens(path, windows, tokenizer): |
| with Path(path).open("w", encoding="utf-8") as f: |
| for window in windows: |
| f.write( |
| f"# window_index={window['window_index']} " |
| f"start={window['start']:.6f} end={window['end']:.6f} " |
| f"prefix_end={window['prefix_end']:.6f} " |
| f"accept=[{window['accept_start']:.6f},{window['accept_end']:.6f}) " |
| f"generation_stop={window['generation_stop']} " |
| f"prefix_tokens={window['prefix_tokens']} " |
| f"tokens={int(window['tokens'].numel())}\n" |
| ) |
| for index, token in enumerate(window["tokens"].tolist()): |
| f.write(f"{index}\t{token}\t{tokenizer.describe(token)}\n") |
| f.write("\n") |
|
|
|
|
| @torch.inference_mode() |
| def generate(model, input_values, prompts=None, *, attention_mask=None, max_length=None, |
| max_new_tokens=None, prefix_tokens=None, |
| autocast_dtype=torch.bfloat16, stop_time_seconds=None, |
| return_dict_in_generate=False, output_logits=False, output_scores=False, |
| progress_callback=None): |
| """Generate symbolic token sequences with the model's event grammar. |
| |
| Inputs are 24 kHz waveforms [batch, samples]. Scores/logits are per generated |
| step, after the prompt prefix, with completed batch items filled by NaNs. |
| """ |
| from transformers.utils import ModelOutput |
| if input_values.ndim == 1: |
| input_values = input_values.unsqueeze(0) |
| if attention_mask is not None and attention_mask.ndim == 1: |
| attention_mask = attention_mask.unsqueeze(0) |
| input_values = input_values.to(next(model.parameters()).device) |
| prompts = model.tokenizer.normalize_prompts(prompts or FULL_TASK_PROMPTS) |
| prefix = list(prefix_tokens) if prefix_tokens is not None else model.tokenizer.prompt_prefix(prompts) |
| prefix_length = len(prefix) - (1 if prefix and prefix[-1] == model.tokenizer.eos_token else 0) |
| if max_new_tokens is not None: |
| if max_length is not None: |
| raise ValueError("Specify max_new_tokens or max_length, not both.") |
| if not isinstance(max_new_tokens, int) or isinstance(max_new_tokens, bool) or max_new_tokens < 1: |
| raise ValueError("max_new_tokens must be a positive integer.") |
| max_length = prefix_length + max_new_tokens |
| limit = model.max_output_seq_len if max_length is None else max_length |
| if not isinstance(limit, int) or isinstance(limit, bool) or not prefix_length < limit <= model.max_output_seq_len: |
| raise ValueError("The token limit must exceed the prefix length and fit the decoder context.") |
| memory = None |
| if attention_mask is not None: |
| with inference_autocast(input_values.device, autocast_dtype): |
| memory = model.get_audio_features(input_values, attention_mask=attention_mask).encoder_last_hidden_state |
| raw, scores, positions = [], [], [] |
| batch_size = input_values.shape[0] |
| def capture(position, ids, logits, masked): |
| positions.append(position) |
| for enabled, values, output in ((output_logits, logits, raw), (output_scores, masked, scores)): |
| if enabled: |
| row = torch.full((batch_size, values.shape[-1]), float('nan'), dtype=torch.float32) |
| row[ids] = values.detach().cpu() |
| output.append(row) |
| tokens = constrained_prompt_generate_batch( |
| model, input_values, prompts, limit, |
| prefix_tokens=prefix_tokens, autocast_dtype=autocast_dtype, |
| stop_time_seconds=stop_time_seconds, progress_callback=progress_callback, |
| step_callback=capture if output_logits or output_scores else None, |
| memory=memory, |
| ) |
| sequences = torch.nn.utils.rnn.pad_sequence(tokens, batch_first=True, padding_value=model.tokenizer.pad_token) |
| if return_dict_in_generate or output_logits or output_scores: |
| return ModelOutput(sequences=sequences, logits=tuple(raw) if output_logits else None, |
| scores=tuple(scores) if output_scores else None, |
| token_positions=torch.tensor(positions, dtype=torch.long)) |
| return sequences |
|
|
