Function _get_unified_codes

moss_tts_delay/llama_cpp/processor.py:167–220

Build the multi-channel (text + audio) packed sequence.

(
    tokenizer: Tokenizer,
    content: str,
    audio_codes_list: list[np.ndarray],
    is_user: bool = True,
    truncation: bool = False,
)

def _get_unified_codes(
    tokenizer: Tokenizer,
    content: str,
    audio_codes_list: list[np.ndarray],
    is_user: bool = True,
    truncation: bool = False,
) -> np.ndarray:
    """Build the multi-channel (text + audio) packed sequence."""
    text_ids = np.array(tokenizer.encode(content), dtype=np.int64)
    n_vq = N_VQ

    if len(audio_codes_list) == 0:
        audio_channel = np.full((len(text_ids), n_vq), AUDIO_PAD_CODE, dtype=np.int64)
        return np.concatenate([text_ids[:, np.newaxis], audio_channel], axis=1)

    audio_start_indices = np.where(text_ids == AUDIO_START_TOKEN_ID)[0]
    audio_end_indices = np.where(text_ids == AUDIO_END_TOKEN_ID)[0]

    if len(audio_start_indices) != len(audio_codes_list) or len(audio_end_indices) != len(audio_codes_list):
        raise ValueError(
            f"Audio markers ({len(audio_start_indices)} starts, {len(audio_end_indices)} ends) "
            f"don't match codes ({len(audio_codes_list)})"
        )

    delay_parts: list[np.ndarray] = []
    prefix_idx = 0

    for start_idx, end_idx, codes in zip(audio_start_indices, audio_end_indices, audio_codes_list):
        start_idx = int(start_idx)
        end_idx = int(end_idx)

        delayed = apply_delay_pattern(codes, AUDIO_PAD_CODE)

        pad_before = np.full(
            (start_idx - prefix_idx + 1, n_vq), AUDIO_PAD_CODE, dtype=np.int64,
        )
        delay_parts.extend([pad_before, delayed])
        prefix_idx = end_idx

    if truncation:
        delay_parts[-1] = delay_parts[-1][:-(n_vq - 1), :]
    else:
        last_end = int(audio_end_indices[-1])
        pad_after = np.full(
            (len(text_ids) - last_end, n_vq), AUDIO_PAD_CODE, dtype=np.int64,
        )
        delay_parts.append(pad_after)

    delay_audio = np.concatenate(delay_parts, axis=0)

    if len(text_ids) != delay_audio.shape[0]:
        text_ids = text_ids[:delay_audio.shape[0]]

    return np.concatenate([text_ids[:, np.newaxis], delay_audio], axis=1)


def parse_generation_output(