Function parse_generation_output

moss_tts_delay/llama_cpp/processor.py:223–262

Parse the generated output into text content and audio codes.

(
    tokenizer: Tokenizer,
    generation_ids: np.ndarray,
    prompt_len: int,
)

def parse_generation_output(
    tokenizer: Tokenizer,
    generation_ids: np.ndarray,
    prompt_len: int,
) -> tuple[str, np.ndarray]:
    """Parse the generated output into text content and audio codes."""
    gen_part = generation_ids[prompt_len:]
    text_channel = gen_part[:, 0].tolist()
    audio_channels = gen_part[:, 1:]

    audio_start_tok = _get_special_token_str(tokenizer, AUDIO_START_TOKEN_ID)
    gen_slot_tok = _get_special_token_str(tokenizer, AUDIO_ASSISTANT_GEN_SLOT_TOKEN_ID)
    delay_slot_tok = _get_special_token_str(tokenizer, AUDIO_ASSISTANT_DELAY_SLOT_TOKEN_ID)
    audio_end_tok = _get_special_token_str(tokenizer, AUDIO_END_TOKEN_ID)

    raw_text = tokenizer.decode(text_channel)

    pattern = re.compile(
        rf"(?:{re.escape(audio_start_tok)})?"
        rf"(?:{re.escape(gen_slot_tok)})*"
        rf"(?:{re.escape(delay_slot_tok)})*"
        rf"{re.escape(audio_end_tok)}"
    )

    def repl(match: re.Match) -> str:
        seg = match.group(0)
        if gen_slot_tok in seg:
            return AUDIO_PLACEHOLDER
        return ""

    text = pattern.sub(repl, raw_text)

    segments = extract_audio_segments(audio_channels)

    if segments:
        audio_codes = np.concatenate(segments, axis=0)
    else:
        audio_codes = np.zeros((0, N_VQ), dtype=np.int64)

    return text, audio_codes