vllm-omni/vllm_omni/tokenizers/yue2_prompt.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Prompt construction for YuE2-3B, mirroring upstream ``protocol.py``.

The driver (offline example or, later, the serving adapter) builds token-id
prefixes here and submits them as ``prompt_token_ids``. Both generation phases
share one prefix head (instruction + tags + lyrics); the abc phase continues
from ``ABC_START`` and stops at ``ABC_END``, the semantic phase embeds the ABC
score (generated or user-supplied) and continues from ``MUSIC_START``.
"""

from __future__ import annotations

from vllm_omni.model_executor.models.yue2.yue2 import (
    ABC_END,
    ABC_START,
    CODEC_OFFSET,
    EOD,
    MUSIC_END,
    MUSIC_START,
)

INSTRUCTIONS = {
    "off": "Generate music with codec tokens from the given conditions.",
    "melody": (
        "Generate a melody-only ABC transcription without chord symbols, then "
        "generate music with codec tokens from the given conditions."
    ),
    "full": (
        "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions."
    ),
}


def song_text(style: str, lyrics: str, cot: str) -> str:
    return f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n"


def abc_prefix_ids(encode, style: str, lyrics: str, cot: str) -> list[int]:
    """Prompt for the abc phase (cot=full/melody, no external ABC)."""
    if cot == "off":
        raise ValueError("cot=off skips the abc phase entirely")
    return [EOD] + encode(song_text(style, lyrics, cot)) + [ABC_START]


def semantic_prefix_ids(
    encode,
    style: str,
    lyrics: str,
    cot: str,
    abc_ids: list[int] | None = None,
) -> list[int]:
    """Prompt for the semantic phase.

    ``abc_ids`` are the token ids of the (generated or user-supplied) ABC
    score. ``cot=off`` passes ``abc_ids=None`` and skips the span, matching
    upstream ``token_prefixes``.
    """
    base = [EOD] + encode(song_text(style, lyrics, cot))
    if cot == "off":
        return base + [ABC_START, ABC_END, MUSIC_START]
    if abc_ids is None:
        raise ValueError("cot=melody/full requires abc ids")
    if any(not 0 <= t < EOD for t in abc_ids):
        raise ValueError("ABC ids must stay inside the ordinary text vocabulary")
    return base + [ABC_START] + list(abc_ids) + [ABC_END, MUSIC_START]


def abc_ids_from_generated(generated: list[int]) -> list[int]:
    """ABC token ids generated by the abc phase (end token stripped)."""
    ids = list(generated)
    if ids and ids[-1] == ABC_END:
        ids.pop()
    if ABC_START in ids:
        ids = ids[ids.index(ABC_START) + 1 :]
    if any(t >= EOD for t in ids):
        raise ValueError("abc phase produced tokens outside the text vocabulary")
    return ids


def semantic_frames(generated: list[int]) -> list[int]:
    """Codec frame values (0..CODEC_SIZE-1) from semantic-phase output."""

    ids = list(generated)
    if ids and ids[-1] == MUSIC_END:
        ids.pop()
    return [t - CODEC_OFFSET for t in ids]


__all__ = [
    "INSTRUCTIONS",
    "abc_ids_from_generated",
    "abc_prefix_ids",
    "semantic_frames",
    "semantic_prefix_ids",
    "song_text",
]