Function step

moss_tts_delay/llama_cpp/delay_state.py:133–256

Execute one autoregressive step. Returns: next_input_ids: (33,) int64 — [text_token, audio_0, ..., audio_31]

(
    state: DelayState,
    text_logits: np.ndarray,
    audio_logits: np.ndarray,
    config: SamplingConfig,
)

def step(
    state: DelayState,
    text_logits: np.ndarray,
    audio_logits: np.ndarray,
    config: SamplingConfig,
) -> np.ndarray:
    """Execute one autoregressive step.

    Returns:
        next_input_ids: (33,) int64 — [text_token, audio_0, ..., audio_31]
    """
    if state.is_stopping:
        pad_result = np.full(1 + N_VQ, AUDIO_PAD_CODE, dtype=np.int64)
        pad_result[0] = PAD_TOKEN_ID
        return pad_result

    n_vq = N_VQ

    # --- Text token decision ---
    if state.delayed_length < n_vq:
        next_text = AUDIO_ASSISTANT_DELAY_SLOT_TOKEN_ID
    elif state.delayed_length == n_vq:
        next_text = AUDIO_END_TOKEN_ID
        state.is_audio = False
    else:
        text_temp = config.text_temperature if config.text_temperature > 0 else 1.0
        text_do_sample = config.text_temperature > 0
        scaled = text_logits / text_temp

        if not state.is_audio:
            scaled[_PRE_EXCLUDE_IDS] = -np.inf
        else:
            mask = np.ones(scaled.shape[0], dtype=bool)
            mask[_AUDIO_ALLOWED_IDS] = False
            scaled[mask] = -np.inf

        if state.time_step == 0:
            scaled[AUDIO_ASSISTANT_DELAY_SLOT_TOKEN_ID] = -np.inf
        if state.time_step <= n_vq:
            scaled[IM_END_TOKEN_ID] = -np.inf

        next_text = int(sample_token(
            scaled[np.newaxis, :],
            top_p=config.text_top_p,
            top_k=config.text_top_k,
            do_sample=text_do_sample,
        )[0])

    if next_text == AUDIO_START_TOKEN_ID:
        state.is_audio = True
    if next_text == IM_END_TOKEN_ID:
        state.is_stopping = True

    # --- Audio token decision ---
    next_audio = np.full(n_vq, AUDIO_PAD_CODE, dtype=np.int64)

    pre_audio_mask = np.arange(n_vq) < state.audio_length
    if state.delayed_length == INT64_MAX: